YOLOv5-4.0-general.py 源代码导读(通用工具)

YOLOv5介绍

YOLOv5为兼顾速度与性能的目标检测算法。笔者将在近期更新一系列YOLOv5的代码导读博客。YOLOv5为2021.1.5日发布的4.0版本。
YOLOv5开源项目github网址
源代码导读汇总网址
本博客导读的代码为utils文件夹下的general.py,取自1.27日更新的版本。

general.py

该文件提供了模型多个过程中用到的通用方法,每个功能以函数的方式进行定义。

以下为该文件必须导入的模块,其中部分文件来源于其他项目子文件。

# General utils
import glob      # 仅支持部分通配符得文件搜索模块
import logging   # 日志模块
import math      # 数学公式模块
import os        # 与操作系统进行交互的模块
import platform  # 获得操作系统相关信息的模块
import random    # 生成随机数的模块
import re        # 用来匹配字符串(动态、模糊)的模块
import subprocess# 创建子进程的模块
import time      # 用来获取系统时间的模块
from pathlib import Path  #Path对象 简便对path进行操作

import cv2           # opencv库
import numpy as np   # numpy矩阵处理函数库
import torch         # pytorch框架
import torchvision   # 为pytorch 提供一些辅助工具
import yaml          # yaml配置文件模块

# 以下调用三个函数具体注释见 源代码导读的其他文件
from utils.google_utils import gsutil_getsize # 用于返回网站链接对应文件的大小
from utils.metrics import fitness # 返回指标的加权值得行向量
from utils.torch_utils import init_torch_seeds # 功能为初始化随机种子

以下为运行相关的一些基本的设置

# 下两行为设置tensor和numpy array的打印格式 linewidth为每一行字符上限 precision 为精度
torch.set_printoptions(linewidth=320, precision=5, profile='long')
np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format})  # format short g, %precision=5
cv2.setNumThreads(0)  # 阻止opencv参与多线程(与 Pytorch的 Dataloader不兼容)
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8))  # 确定最大的线程数 这里被限制在了8

该函数为对日志的设置进行初始化 rank为-1或0 时设置输出级别为WARN

def set_logging(rank=-1):
    logging.basicConfig(
        format="%(message)s",
        level=logging.INFO if rank in [-1, 0] else logging.WARN)

该函数为初始化随机种子 统一random numpy torch 种子

def init_seeds(seed=0):
    # 初始化随机种子生成器
    random.seed(seed)
    np.random.seed(seed)
    init_torch_seeds(seed)

该函数返回最近的模型 'last.pt’对应的路径

def get_latest_run(search_dir='.'):
    # Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
    # 从Python版本3.5开始,glob模块支持该"**"指令(仅当传递recursive标志时才会解析该指令
    # glob.glob函数匹配所有的符合条件的文件,并将其以list的形式返回
    last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
    # os.getctime 返回路径对应文件的创建时间 
    # 也就是返回所有文件中创建时间最晚的路径
    return max(last_list, key=os.path.getctime) if last_list else ''

用socket模块 检查当前主机网络连接是否可用

def check_online():
    # Check internet connectivity
    import socket
    try:
        socket.create_connection(("1.1.1.1", 53))  # check host accesability 该单词拼错 应为accessbility
        return True
    except OSError:
        return False

检查当前代码是否是最新版 如果不是最新的 会提示使用git pull命令进行升级

def check_git_status():
    # Recommend 'git pull' if code is out of date
    # 彩色显示github单词 colorstr()函数后续介绍
    print(colorstr('github: '), end='')
    try:
        # 检查以git结尾的路径存在
        assert Path('.git').exists(), 'skipping check (not a git repository)'
        # 但是包含"/workspace"的路径不存在
        assert not Path('/workspace').exists(), 'skipping check (Docker image)'  # not Path('/.dockerenv').exists()
        # 保证主机网络是可用的
        assert check_online(), 'skipping check (offline)'
        # 这里是创建cmd命令 并创建子进程进行执行
        cmd = 'git fetch && git config --get remote.origin.url'
        url = subprocess.check_output(cmd, shell=True).decode().strip().rstrip('.git')  # github repo url
        branch = subprocess.check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip()  # checked out
        n = int(subprocess.check_output(f'git rev-list {branch}..origin/master --count', shell=True))  # commits behind
        # n大于0说明 当前版本之后还有commit 因此当前版本不是最新的 s为输出的相关提示
        if n > 0:
            s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \
                f"Use 'git pull' to update or 'git clone {url}' to download latest."
        else:
            s = f'up to date with {url} ✅'
        # 通过.encode().decode()的组合忽略掉无法用ascii编码的内容
        print(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s)
    except Exception as e:
        print(e)

用于检查已经安装的包是否满足file对应txt文件的要求

def check_requirements(file='requirements.txt', exclude=()):

    import pkg_resources #此为管理安装包信息相关模块
    # x在被读取的时候转换为pkg_resources.Requirement类
    # x.name 返回安装包对应名称 x.specifier 返回符号后面内容 即安装包对应版本
    # requirements 负责将上述pkg类转换为列表对象 以供require函数使用
    requirements = [f'{x.name}{x.specifier}' for x in pkg_resources.parse_requirements(Path(file).open())
                    if x.name not in exclude]
    pkg_resources.require(requirements)  #判断已经安装的包是否符合要求

检查img_size是否能被s整除,这里默认s为32

def check_img_size(img_size, s=32):
    # make_divisible函数见下文注释
    new_size = make_divisible(img_size, int(s))  # 返回大于等于img_size的最小能被s整除的值
    if new_size != img_size: #新size和旧size不同时 打印出信息 
        print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
    return new_size # 返回能被s整除的new_size

检查相关文件路径是否能找到文件,如果不能则在全局路径中寻找,匹配到多个时返回第一个

def check_file(file):

    if os.path.isfile(file) or file == '':
        return file
    else:
        files = glob.glob('./**/' + file, recursive=True)  # 寻找文件
        assert len(files), 'File Not Found: %s' % file  # assert 文件没有找到
        assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % (file, files)  # assert unique
        return files[0]  # 返回文件

检查数据集 如果本地没有则从torch库中下载并解压数据集

#          dict内容需要实际调试(暂无法逐行调试) 该方法只能适用某几个数据集
def check_dataset(dict):

    val, s = dict.get('val'), dict.get('download')
    if val and len(val):
        val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])]  # val path
        if not all(x.exists() for x in val):
            print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
            if s and len(s):  # download script
                print('Downloading %s ...' % s)
                if s.startswith('http') and s.endswith('.zip'):  # URL
                    f = Path(s).name  # filename
                    torch.hub.download_url_to_file(s, f)
                    r = os.system('unzip -q %s -d ../ && rm %s' % (f, f))  # unzip
                else:  # bash script
                    r = os.system(s)
                print('Dataset autodownload %s\n' % ('success' if r == 0 else 'failure'))  # analyze return value
            else:
                raise Exception('Dataset not found.')

取大于等于x的最小值,该值能被divisor整除

def make_divisible(x, divisor):
    # 如 100,3 返回34 100,5返回20
    return math.ceil(x / divisor) * divisor

对字符串s里在pattern中字符替换为下划线_ pattern中[]不能省

def clean_str(s):
    
    return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)

定义从y1增加到y2 的余弦上升函数

# x取值从0-steps对应y1-y2 之后进行循环
def one_cycle(y1=0.0, y2=1.0, steps=100):

    return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1

在命令行中,将文字带有颜色进行输出

def colorstr(*input):
    # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e.  colorstr('blue', 'hello world')
    # args为输入的颜色序列,string为输入的字符串 eles后内容为默认设置 colors字典为颜色对应的字符表示 该表示只在命令行有效
    *args, string = input if len(input) > 1 else ('blue', 'bold', input[0])  # color arguments, string
    colors = {'black': '\033[30m',  # basic colors
              'red': '\033[31m',
              'green': '\033[32m',
              'yellow': '\033[33m',
              'blue': '\033[34m',
              'magenta': '\033[35m',
              'cyan': '\033[36m',
              'white': '\033[37m',
              'bright_black': '\033[90m',  # bright colors
              'bright_red': '\033[91m',
              'bright_green': '\033[92m',
              'bright_yellow': '\033[93m',
              'bright_blue': '\033[94m',
              'bright_magenta': '\033[95m',
              'bright_cyan': '\033[96m',
              'bright_white': '\033[97m',
              'end': '\033[0m',  # misc
              'bold': '\033[1m',
              'underline': '\033[4m'}
    # 把一个string的开头和末尾加上颜色 命令行输出会更好看
    return ''.join(colors[x] for x in args) + f'{string}' + colors['end']

从训练标签获得类权重(反频率)

"""
  labels Array(M,N,5) [class x y w h] 猜测因为函数内有concatenate函数降维
  从训练标签获得类权重(反频率) 
  weights最终权重为 每一类出现次数的倒数 占全部类别的百分比(出现次数为0是认为 为1)
  weights实际的物理意义为:次数出现越少的越重要 没有出现的最重要 对应weights最大
"""
def labels_to_class_weights(labels, nc=80):
    
    if labels[0] is None:  # 如果labels为空 返回一个空tensor
        return torch.Tensor()
    # np.concatenate(array,0) 为将数组array的高维元素合并 转换为低一维的元素
    labels = np.concatenate(labels, 0)  # labels.shape = (866643, 5) for COCO
    # classes Array(1,N) 代表labels的类别
    classes = labels[:, 0].astype(np.int)  # labels = [class x y w h]
    # weights Array(1,N) 从第0类到第nc类 每一类在classes中出现的次数
    weights = np.bincount(classes, minlength=nc)  # 每一种类别出现的次数

    # Prepend gridpoint count (for uCE training)
    # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum()  # gridpoints per image
    # weights = np.hstack([gpi * len(labels)  - weights.sum() * 9, weights * 9]) ** 0.5  # prepend gridpoints to start

    weights[weights == 0] = 1  # 将出现次数为0的类别 统一替换为1
    weights = 1 / weights      # weights改为 类别出现次数的倒数
    weights /= weights.sum()   # 每一个元素占总和的百分比
    # torch.from_numpy()函数从numpy数组构建tensor 两者共享内存
    return torch.from_numpy(weights)

通过图像的labels信息返回图像对应的权重

# 由label计算图像的权重
def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
    
    class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels])
    # 实际计算中class_weights 应由labels_to_class_weights计算得到
    image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
    # index = random.choices(range(n), weights=image_weights, k=1)  # weight image sample
    
    return image_weights

将80类的coco索引转换为91类的coco索引 x为80类中的每一类 在91类中对应的位置

def coco80_to_coco91_class():  # converts 80-index (val2014) to 91-index (paper)
    # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
    # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
    # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
    # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)]  # darknet to coco
    # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)]  # coco to darknet
    x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
         35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
         64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
    return x

以下几个函数 x Tensor/Array(N,4) 为转换不同的坐标格式

# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
def xyxy2xywh(x):
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center 
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y

# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
def xywh2xyxy(x):
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y

# Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
def xywhn2xyxy(x, w=640, h=640, padw=32, padh=32): #xywh 为相关比例 所以转换的时候需要乘图像的尺寸
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw  # top left x1
    y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh  # top left y1
    y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw  # bottom right x2
    y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh  # bottom right y2
    return y

将坐标按照不同图像比例进行放缩 从img1的尺寸方放缩到img0的尺寸
pad为扩充的间隔值 参考CNN中的padding

def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):

    if ratio_pad is None:  # calculate from img0_shape
        # gain 取长宽中的较小值 old/new 小 说明 new/old 大  则gain为长或宽中扩大比例较大的那个
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
        # pad = a,b 格式等价于 pad = (a,b)
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
    else:
        gain = ratio_pad[0][0] #指定比例
        pad = ratio_pad[1]     #指定pad值

    coords[:, [0, 2]] -= pad[0]  # x padding
    coords[:, [1, 3]] -= pad[1]  # y padding
    coords[:, :4] /= gain # /gain = * new/old
    clip_coords(coords, img0_shape) #见下函数解析
    return coords

将bounding box坐标范围限制在图像尺寸中

# torch.clamp_(min,max) 函数能够将tensor的取值限制在(min,max)之间,超出这个范围自动归到这个范围中
# 将bounding box坐标范围限制在图像尺寸中
def clip_coords(boxes, img_shape):

    boxes[:, 0].clamp_(0, img_shape[1])  # x1
    boxes[:, 1].clamp_(0, img_shape[0])  # y1
    boxes[:, 2].clamp_(0, img_shape[1])  # x2
    boxes[:, 3].clamp_(0, img_shape[0])  # y2

返回box1 到 box2的IOU值. box1 is 4, box2 is nx4

# 这里 box1是一个label box2是N个label
def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9):

    box2 = box2.T

    # Get the coordinates of bounding boxes
    if x1y1x2y2:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
    else:  # xywh 到 xyxy 的坐标转换 因为广播机制 b2系列均为N个元素的一维数组
        b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
        b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
        b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
        b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2

    # Intersection area 计算相交的区域 (可用一维数轴理解计算过程)
    # 相交的区域一定是一个矩形 所以计算相交的x和y的长度 两者等价 下以x距离
    # 右下顶点代表了x较大的取值 对两者右下顶点的x取最小 代表了box向外界的最小值
    # 左上顶点......较小............左上顶点......大  .........内界的最大值
    # 当外最小减内最大 为 负数时 代表两者没有交集 所以调用.clamp()函数直接将iou置零
    inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
            (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)

    # Union Area 计算相并的区域 总面积 减去多余计算的相交的面积
    w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
    w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
    union = w1 * h1 + w2 * h2 - inter + eps

    iou = inter / union # 交/并
    # 以下为GIOU DIOU CIOU 的计算过程
    if GIoU or DIoU or CIoU:
        cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex (smallest enclosing box) width
        ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
        if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
                    (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center distance squared
            if DIoU:
                return iou - rho2 / c2  # DIoU
            elif CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
                with torch.no_grad():
                    alpha = v / ((1 + eps) - iou + v)
                return iou - (rho2 / c2 + v * alpha)  # CIoU
        else:  # GIoU https://arxiv.org/pdf/1902.09630.pdf
            c_area = cw * ch + eps  # convex area
            return iou - (c_area - union) / c_area  # GIoU
    else:
        return iou  # IoU

返回iou Tensor[N,M] 其中(n,m)坐标代表 box1中的第n个框和box2中的第m个框的iou值

def box_iou(box1, box2):
    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
    """
    Return intersection-over-union (Jaccard index) of boxes.
    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
    Arguments:
        box1 (Tensor[N, 4])
        box2 (Tensor[M, 4])
    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise
            IoU values for every element in boxes1 and boxes2
    """

    def box_area(box):
        # box = 4xn
        return (box[2] - box[0]) * (box[3] - box[1])

    area1 = box_area(box1.T)
    area2 = box_area(box2.T)
     
    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
    # 相交面积N*M矩阵的每一个元素 为 (右下-左上)取大于0  在对第三维做乘积
    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
    return inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)

返回根据wh点确定的iou值

def wh_iou(wh1, wh2):
    # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
    wh1 = wh1[:, None]  # [N,1,2]
    wh2 = wh2[None]  # [1,M,2]
    inter = torch.min(wh1, wh2).prod(2)  # [N,M]
    return inter / (wh1.prod(2) + wh2.prod(2) - inter)  # iou = inter / (area1 + area2 - inter)

该函数为非极大值抑制算法(NMS) 用于筛选bounding box

def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):
    """Performs Non-Maximum Suppression (NMS) on inference results
    Returns:
         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
    """
    # prediction是三维:包含了若干张图片
    # prediction.shape[0]代表的是图片的数目 
    # prediction.shape[1]代表的是每一张图片的每一行box信息 shape[2]代表的是列数
    
    nc = prediction.shape[2] - 5  # number of classes 
    xc = prediction[..., 4] > conf_thres  # 置信度 x是布尔数组

    # Settings
    min_wh, max_wh = 2, 4096  # box长宽的最小最大像素值
    max_det = 300  # 每张照片检测的最大数量
    max_nms = 30000  # 在torchvision.ops.nms()中的box最大数目
    time_limit = 10.0  # 几秒后退出
    redundant = True  # 是否需要redundant detections
    multi_label = nc > 1  # 每个盒子有多个标签 (adds 0.5ms/img) 类别大于1?
    merge = False  # 是否使用 merge-NMS

    t = time.time() #t为当前的时间
    # 创建了一个大list 其中每一个元素 都是一个空的(0,6) tensor 该output最终作为结果输出
    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # xi:image index, x:image inference
        # Apply constraints
        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
        # x Tensor(center x, center y, width, height, cls_conf, obj_conf)
        x = x[xc[xi]]  # 依据置信度筛选了图像bounding box的信息

        # 获得先验标签 如果自动标记的话
        if labels and len(labels[xi]):
            l = labels[xi]
            v = torch.zeros((len(l), nc + 5), device=x.device)
            v[:, :4] = l[:, 1:5]  # box
            v[:, 4] = 1.0  # conf
            v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
            x = torch.cat((x, v), 0) #这里是将自动标注的标签 添加到x之中 一起进行NMS

        # 如果该张图片没有坐标了 则进行下一张图片
        if not x.shape[0]:
            continue

        # 计算置信度 将cls_conf列置为 conf*cls的乘积
        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

        # 格式转换 Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])

        # 建立检测矩阵 nx6 (xyxy, conf, cls)
        # torch.cat()为tensor的拼接操作
        if multi_label:
            # torch.nonzero()函数返回非零值对应的值 以及对应的索引
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
        else:  # 只取一类
            conf, j = x[:, 5:].max(1, keepdim=True)
            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]

        # 根据定义的classes 进行筛选
        if classes is not None:
            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]

        # Apply finite constraint
        # if not torch.isfinite(x).all():
        #     x = x[torch.isfinite(x).all(1)]

        # 检查shape (box的数量)
        n = x.shape[0]  # number of boxes
        if not n:  # 如果box 的数量为0 检查下一张图片
            continue
        elif n > max_nms:  # 超过了处理的数目 则根据置信度排列并取到max_nms张图片
            x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence

        # Batched NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes 根据类进行整理(offset by class), scores
        # 直接调用了pytorch的库进行NMS 操作
        i = torchvision.ops.nms(boxes, scores, iou_thres)  
        if i.shape[0] > max_det:  # 限定检测数量
            i = i[:max_det]
        # Merge NMS 在hard-nms的基础上,增加保留框位置平滑策略(重叠框位置信息求解平均值),使框的位置更加精确
        if merge and (1 < n < 3E3):  # 采用Merge NMS (boxes merged using weighted mean)
            # 更新 boxes: boxes(i,4) = weights(i,n) * boxes(n,4)
            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
            weights = iou * scores[None]  # box weights
            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
            if redundant:
                i = i[iou.sum(1) > 1]  # require redundancy

        output[xi] = x[i]
        # 如果超过了NMS的处理时间上限则进行报错
        if (time.time() - t) > time_limit:
            print(f'WARNING: NMS time limit {time_limit}s exceeded')
            break  # time limit exceeded

    return output

将模型从优化器剥离以完成训练 或可保存为s

def strip_optimizer(f='weights/best.pt', s=''):  # from utils.general import *; strip_optimizer()
    
    # x为加载训练的模型
    x = torch.load(f, map_location=torch.device('cpu'))
    #以下为将 模型训练涉及到的若干个指定变量置空
    for key in 'optimizer', 'training_results', 'wandb_id': 
        x[key] = None
    x['epoch'] = -1
    x['model'].half()  # 转换为FP16精度
    for p in x['model'].parameters():
        p.requires_grad = False
    #保存模型
    torch.save(x, s or f)
    mb = os.path.getsize(s or f) / 1E6  # 获得filesize 并转换为MB
    print('Optimizer stripped from %s,%s %.1fMB' % (f, (' saved as %s,' % s) if s else '', mb))

打印超参进化的结果在evolve.txt (使用这个功能需要 train.py --evolve)

def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''):
    
    # 定义相关变量 并赋值 按指定格式输出
    a = '%10s' * len(hyp) % tuple(hyp.keys())  # hyperparam keys
    b = '%10.3g' * len(hyp) % tuple(hyp.values())  # hyperparam values
    c = '%10.4g' * len(results) % results  # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
    print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c))
    
    # 如果大于本地,则下载evolve.txt
    if bucket:
        url = 'gs://%s/evolve.txt' % bucket
        if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0):
            os.system('gsutil cp %s .' % url) 
    # 打开ecolve.txt 并添加结果
    with open('evolve.txt', 'a') as f:  
        f.write(c + b + '\n')
    x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0)  # load unique rows
    x = x[np.argsort(-fitness(x))]  # sort
    np.savetxt('evolve.txt', x, '%10.3g')  # 根据fitness整理后的结果保存

    # 保存yaml配置文件 为'hyp_evolved.yaml'
    for i, k in enumerate(hyp.keys()):
        hyp[k] = float(x[0, i + 7])
    with open(yaml_file, 'w') as f:
        results = tuple(x[0, :7])
        c = '%10.4g' * len(results) % results  # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
        f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n')
        yaml.dump(hyp, f, sort_keys=False)

    if bucket:
        os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket))  # upload

定义了一个二级分类器来处理yolo的输出

def apply_classifier(x, model, img, im0):

    im0 = [im0] if isinstance(im0, np.ndarray) else im0
    for i, d in enumerate(x):  # i为index d代表每一张图片包含的信息
        if d is not None and len(d):
            d = d.clone()

            # Reshape and pad cutouts
            b = xyxy2xywh(d[:, :4])  # b为xywh格式的bounding box
            b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1)  # rectangle to square
            b[:, 2:] = b[:, 2:] * 1.3 + 30  # pad 操作 1.3和30看起来像是随便加的
            d[:, :4] = xywh2xyxy(b).long() #将d前四列坐标格式转换成xyxy格式

            # 将box根据图像size的不同进行rescale 该函数上文已经注释
            scale_coords(img.shape[2:], d[:, :4], im0[i].shape)

            # pred_cls1获得每一个预测值对应的类别
            pred_cls1 = d[:, 5].long()
            
            # ims汇总图片中所有子目标的区域
            ims = []
            for j, a in enumerate(d):  # per item
                # 根据xyxy格式的标注对原始图像进行剪裁
                cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
                # 用opencv进行resize操作
                im = cv2.resize(cutout, (224, 224))  # opencv图像储存方式为BGR
                # cv2.imwrite('test%i.jpg' % j, cutout)
                
                # 将BGR格式的图像转换为RGB (3*416*416)
                im = im[:, :, ::-1].transpose(2, 0, 1)  
                # 将img从uint8格式转换为float 32 格式
                im = np.ascontiguousarray(im, dtype=np.float32)  
                im /= 255.0  # 0 - 255 to 0.0 - 1.0
                ims.append(im) #添加剪裁后图像的信息
            
            # 用model模型进行预测
            pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1)  
            # 保留预测一致的结果
            x[i] = x[i][pred_cls1 == pred_cls2]  

    return x

递增路径如 runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc.

# 如:将训练结果统一保存到runs/train/exp文件夹中
#    如果该文件夹已经存在 则将路径修改为 runs/train/exp1
def increment_path(path, exist_ok=True, sep=''):

    path = Path(path)  # os-agnostic
    if (path.exists() and exist_ok) or (not path.exists()):
        return str(path)
    else:
        dirs = glob.glob(f"{path}{sep}*")  # 相似的路径
        matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
        i = [int(m.groups()[0]) for m in matches if m]  # 索引
        n = max(i) + 1 if i else 2  # 递增数字
        return f"{path}{sep}{n}"  # 更新路径

水平有限,不当之处欢迎读者补充!

已标记关键词 清除标记
相关推荐
©️2020 CSDN 皮肤主题: 大白 设计师:CSDN官方博客 返回首页