【YOLOV5-5.x 源码解读】val.py

前言

源码： YOLOv5源码.
导航: 【YOLOV5-5.x 源码讲解】整体项目文件导航.
注释版全部项目文件已上传至GitHub: yolov5-5.x-annotations.

这个文件主要是在每一轮训练结束后，验证当前模型的mAP、混淆矩阵等指标。这个文件本来是叫test.py的，但我觉得叫val.py可能更加恰当些，所以改了。

实际上这个脚本最常用的应该是通过train.py调用 run 函数，而不是通过执行 val.py 的。所以在了解这个脚本的时候，其实最重要的就是 run 函数。

0、导入需要的包和基本配置

import argparse                  # 解析命令行参数模块
import json                      # 实现字典列表和JSON字符串之间的相互解析
import os                        # 与操作系统进行交互的模块 包含文件路径操作和解析
import sys                       # sys系统模块 包含了与Python解释器和它的环境有关的函数
from pathlib import Path         # Path将str转换为Path对象 使字符串路径易于操作的模块
from threading import Thread     # 线程操作模块import numpy as np
import torch
import yaml
from tqdm import tqdmFILE = Path(__file__).absolute()  # FILE = WindowsPath 'F:\yolo_v5\yolov5-U\val.py'
# 将'F:/yolo_v5/yolov5-U'加入系统的环境变量 该脚本结束后失效
sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to pathfrom models.experimental import attempt_load
from utils.datasets import create_dataloader
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
from utils.metrics import ap_per_class, ConfusionMatrix
from utils.plots import plot_images, output_to_target, plot_study_txt
from utils.torch_utils import select_device, time_synchronized

1、设置opt参数

def parse_opt():"""opt参数详解data: 数据集配置文件地址 包含数据集的路径、类别个数、类名、下载地址等信息weights: 模型的权重文件地址 weights/yolov5s.ptbatch_size: 前向传播的批次大小 默认32imgsz: 输入网络的图片分辨率 默认640conf-thres: object置信度阈值 默认0.001iou-thres: 进行NMS时IOU的阈值 默认0.6task: 设置测试的类型 有train, val, test, speed or study几种 默认valdevice: 测试的设备single-cls: 数据集是否只用一个类别 默认Falseaugment: 测试是否使用TTA Test Time Augment 默认Falseverbose: 是否打印出每个类别的mAP 默认False下面三个参数是auto-labelling(有点像RNN中的teaching forcing)相关参数详见:https://github.com/ultralytics/yolov5/issues/1563 下面解释是作者原话save-txt: traditional auto-labellingsave-hybrid: save hybrid autolabels, combining existing labels with new predictions before NMS (existing predictions given confidence=1.0 before NMS.save-conf: add confidences to any of the above commandssave-json: 是否按照coco的json格式保存预测框，并且使用cocoapi做评估（需要同样coco的json格式的标签） 默认Falseproject: 测试保存的源文件 默认runs/testname: 测试保存的文件地址 默认exp 保存在runs/test/exp下exist-ok: 是否存在当前文件 默认False 一般是 no exist-ok 连用 所以一般都要重新创建文件夹half: 是否使用半精度推理 默认False"""parser = argparse.ArgumentParser(prog='val.py')parser.add_argument('--data', type=str, default='data/coco128.yaml', help='dataset.yaml path')parser.add_argument('--weights', nargs='+', type=str, default='weights/yolov5s.pt', help='model.pt path(s)')parser.add_argument('--batch-size', type=int, default=4, help='batch size')parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')parser.add_argument('--task', default='val', help='train, val, test, speed or study')parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')parser.add_argument('--augment', action='store_true', help='augmented inference')parser.add_argument('--verbose', action='store_true', help='report mAP by class')parser.add_argument('--save-txt', default=True, action='store_true', help='save results to *.txt')parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')parser.add_argument('--save-conf', default=True, action='store_true', help='save confidences in --save-txt labels')parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')parser.add_argument('--project', default='runs/test', help='save to project/name')parser.add_argument('--name', default='exp', help='save to project/name')parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')opt = parser.parse_args()   # 解析上述参数opt.save_json |= opt.data.endswith('coco.yaml')  # |或 左右两个变量有一个为True 左边变量就为Trueopt.save_txt |= opt.save_hybridopt.data = check_file(opt.data)  # check filereturn opt

2、执行main函数

def main(opt):# 初始化日志 设置logging的等级level、输出格式和内容format等基本信息set_logging()# 打印信息 test(高亮): opt各个参数值print(colorstr('test: ') + ', '.join(f'{
      k}={
      v}' for k, v in vars(opt).items()))# 检测requirements文件中需要的包是否安装好了check_requirements(exclude=('tensorboard', 'thop'))# 如果task in ['train', 'val', 'test']就正常测试 训练集/验证集/测试集if opt.task in ('train', 'val', 'test'):  # run normallyrun(**vars(opt))# 如果opt.task == 'speed' 就测试yolov5系列和yolov3-spp各个模型的速度评估elif opt.task == 'speed':  # speed benchmarksfor w in opt.weights if isinstance(opt.weights, list) else [opt.weights]:run(opt.data, weights=w, batch_size=opt.batch_size, imgsz=opt.imgsz, conf_thres=.25, iou_thres=.45,save_json=False, plots=False)# 如果opt.task = ['study']就评估yolov5系列和yolov3-spp各个模型在各个尺度下的指标并可视化elif opt.task == 'study':  # run over a range of settings and save/plot# python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5s.pt yolov5m.pt yolov5l.pt yolov5x.ptx = list(range(256, 1536 + 128, 128))  # x axis (image sizes)for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]:f = f'study_{
      Path(opt.data).stem}_{
      Path(w).stem}.txt'  # filename to save toy = []  # y axisfor i in x:  # img-sizeprint(f'\nRunning {
      f} point {
      i}...')r, _, t = run(opt.data, weights=w, batch_size=opt.batch_size, imgsz=i, conf_thres=opt.conf_thres,iou_thres=opt.iou_thres, save_json=opt.save_json, plots=False)y.append(r + t)  # results and timesnp.savetxt(f, y, fmt='%10.4g')  # saveos.system('zip -r study.zip study_*.txt')# 可视化各个指标plot_study_txt(x=x)  # plot

可以看到这个模块根据opt.task可以分为三个分支，我们主要的分支还是在 opt.task in (‘train’, ‘val’, ‘test’)。而其他的两个分支，大家大概看看在干什么就可以了，没什么用。一般我们都是直接进入第一个分支，执行run函数。

3、执行run函数

run 函数其实用train.py执行的，并不是执行val.py。

train.py调用（每个训练epoch后验证当前模型）:

3.1、载入参数

@torch.no_grad()      # 不参与反向传播
def run(data, weights=None, batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.6,task='val', device='', single_cls=False, augment=False, verbose=False, save_txt=False,save_hybrid=False, save_conf=False, save_json=False, project='runs/test', name='exp',exist_ok=False, half=True, model=None, dataloader=None, save_dir=Path(''), plots=True,wandb_logger=None, compute_loss=None,):""":params data: 数据集配置文件地址 包含数据集的路径、类别个数、类名、下载地址等信息 train.py时传入data_dict:params weights: 模型的权重文件地址 运行train.py=None 运行test.py=默认weights/yolov5s.pt:params batch_size: 前向传播的批次大小 运行test.py传入默认32 运行train.py则传入batch_size // WORLD_SIZE * 2:params imgsz: 输入网络的图片分辨率 运行test.py传入默认640 运行train.py则传入imgsz_test:params conf_thres: object置信度阈值 默认0.25:params iou_thres: 进行NMS时IOU的阈值 默认0.6:params task: 设置测试的类型 有train, val, test, speed or study几种 默认val:params device: 测试的设备:params single_cls: 数据集是否只用一个类别 运行test.py传入默认False 运行train.py则传入single_cls:params augment: 测试是否使用TTA Test Time Augment 默认False:params verbose: 是否打印出每个类别的mAP 运行test.py传入默认Fasle 运行train.py则传入nc < 50 and final_epoch:params save_txt: 是否以txt文件的形式保存模型预测框的坐标 默认True:params save_hybrid: 是否save label+prediction hybrid results to *.txt 默认False:params save_conf: 是否保存预测每个目标的置信度到预测tx文件中 默认True:params save_json: 是否按照coco的json格式保存预测框，并且使用cocoapi做评估（需要同样coco的json格式的标签）运行test.py传入默认Fasle 运行train.py则传入is_coco and final_epoch(一般也是False):params project: 测试保存的源文件 默认runs/test:params name: 测试保存的文件地址 默认exp 保存在runs/test/exp下:params exist_ok: 是否存在当前文件 默认False 一般是 no exist-ok 连用 所以一般都要重新创建文件夹:params half: 是否使用半精度推理 FP16 half-precision inference 默认False:params model: 模型 如果执行test.py就为None 如果执行train.py就会传入ema.ema(ema模型):params dataloader: 数据加载器 如果执行test.py就为None 如果执行train.py就会传入testloader:params save_dir: 文件保存路径 如果执行test.py就为‘’ 如果执行train.py就会传入save_dir(runs/train/expn):params plots: 是否可视化 运行test.py传入默认True 运行train.py则传入plots and final_epoch:params wandb_logger: 网页可视化 类似于tensorboard 运行test.py传入默认None 运行train.py则传入wandb_logger(train):params compute_loss: 损失函数 运行test.py传入默认None 运行train.py则传入compute_loss(train):return (Precision, Recall, map@0.5, map@0.5:0.95, box_loss, obj_loss, cls_loss)"""

3.2、初始化配置1

训练时（train.py）调用：初始化模型参数、训练设备
验证时（val.py）调用：初始化设备、save_dir文件路径、make dir、加载模型、check imgsz、 加载+check data配置信息

    # ============================================== 1、初始化配置1 ==================================================# 初始化模型并选择相应的计算设备# 判断是否是训练时调用run函数(执行train.py脚本), 如果是就使用训练时的设备 一般都是traintraining = model is not Noneif training:  # called by train.pydevice = next(model.parameters()).device  # get model device# 如果不是trin.py调用run函数(执行val.py脚本)就调用select_device选择可用的设备# 并生成save_dir + make dir + 加载模型model + check imgsz + 加载data配置信息else:device = select_device(device, batch_size=batch_size)# 生成save_dir文件路径 run\test\expnsave_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run# make dir run\test\expn\labels(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)# 加载模型 load FP32 model 只在运行test.py才需要自己加载modelmodel = attempt_load(weights, map_location=device)# gs: 模型最大的下采样stride 一般[8, 16, 32] 所有gs一般是32gs = max(int(model.stride.max()), 32)  # grid size (max stride)# 检测输入图片的分辨率imgsz是否能被gs整除 只在运行test.py才需要自己生成check imgsz# imgsz_testimgsz = check_img_size(imgsz, s=gs)  # check image size# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99# if device.type != 'cpu' and torch.cuda.device_count() > 1:# model = nn.DataParallel(model)# Data 加载数据配置信息 只有运行test.py才需要加载数据配置信息, 因为它需要根据data生成新的dataloader# 而运行train.py时是直接传入testloader的, 所以不需要加载数据配置信息with open(data, encoding='utf-8') as f:data = yaml.safe_load(f)check_dataset(data)  # check

3.3、调整模型

半精度验证half model + 模型剪枝prune + 模型融合conv+bn

    # ============================================== 2、调整模型设置 ==================================================# Half model 只能在单GPU设备上才能使用# 一旦使用half, 不但模型需要设为half, 输入模型的图片也需要设为halfhalf &= device.type != 'cpu'  # half precision only supported on CUDAif half:model.half()# from utils.torch_utils import prune# prune(model, 0.3) # 模型剪枝# model = model.fuse() # 模型融合 融合conv+bnmodel.eval()  # 启动模型验证模式

3.4、初始化配置2

是否是coco数据集is_coco + 类别个数nc + 计算mAP相关参数 + 初始化日志 Logging

    # ============================================== 3、初始化配置2 ==================================================# 测试数据是否是coco数据集is_coco = type(data['val']) is str and data['val'].endswith('coco/val2017.txt')  # COCO dataset boolnc = 1 if single_cls else int(data['nc'])  # number of classes# 计算mAP相关参数# 设置iou阈值 从0.5-0.95取10个(0.05间隔) iou vector for mAP@0.5:0.95# # iouv: [0.50000, 0.55000, 0.60000, 0.65000, 0.70000, 0.75000, 0.80000, 0.85000, 0.90000, 0.95000]iouv = torch.linspace(0.5, 0.95, 10).to(device)# mAP@0.5:0.95 iou个数=10个niou = iouv.numel()# 初始化日志 Logginglog_imgs = 0if wandb_logger and wandb_logger.wandb:log_imgs = min(wandb_logger.log_imgs, 100)

3.5、加载val数据集

训练时（train.py）调用：加载val数据集
验证时（val.py）调用：不需要加载val数据集 直接从train.py 中传入testloader

    # ============================================== 4、加载val数据集 ==================================================# 如果不是训练(执行val.py脚本调用run函数)就调用create_dataloader生成dataloader# 如果是训练(执行train.py调用run函数)就不需要生成dataloader 可以直接从参数中传过来testloaderif not training:if device.type != 'cpu':# 这里创建一个全零数组测试下前向传播是否能够正常运行model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run oncetask = task if task in ('train', 'val', 'test') else 'val'  # path to train/val/test images# 创建dataloader 这里的rect默认为True 矩形推理用于测试集 在不影响mAP的情况下可以大大提升推理速度dataloader = create_dataloader(data[task], imgsz, batch_size, gs, single_cls, pad=0.5, rect=True,prefix=colorstr(f'{
      task}: '))[0]

3.6、初始化配置3

初始化混淆矩阵 + 数据集类名 + 获取coco数据集的类别索引 + 设置tqdm进度条 + 初始化p, r, f1, mp, mr, map50, map指标和时间t0, t1, t2 + 初始化测试集的损失 + 初始化json文件中的字典 统计信息 ap等

    # ============================================== 5、初始化配置3 ==================================================# 初始化一些测试需要的参数seen = 0  # 初始化测试的图片的数量# 初始化混淆矩阵confusion_matrix = ConfusionMatrix(nc=nc)# 获取数据集所有类别的类名names = {
    k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}# 获取coco数据集的类别索引# coco数据集是80个类 索引范围本应该是0~79,但是这里返回的确是0~90 coco官方就是这样规定的# coco80_to_coco91_class就是为了与上述索引对应起来，返回一个范围在0~80的索引数组coco91class = coco80_to_coco91_class()# 设置tqdm进度条的显示信息s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')# 初始化p, r, f1, mp, mr, map50, map指标和时间t0, t1, t2p, r, f1, mp, mr, map50, map, t0, t1, t2 = 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.# 初始化测试集的损失loss = torch.zeros(3, device=device)# 初始化json文件中的字典 统计信息 ap等jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []

3.7、开始验证

    # ============================================== 6、开始验证 ==================================================for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):

3.7.1、预处理图片和target

        # 6.1、预处理图片和targett_ = time_synchronized()  # 获取当前时间img = img.to(device, non_blocking=True)  # img to device# 如果half为True 就把图片变为half精度 uint8 to fp16/32img = img.half() if half else img.float()img /= 255.0  # 0 - 255 to 0.0 - 1.0targets = targets.to(device)  # targets to device# batch size, channels, height, widthnb, _, height, width = img.shapet = time_synchronized()  # 获取当前时间t0 += t - t_  # t0: 累计处理数据时间

3.7.2、model 前向推理

        # 6.2、Run model 前向推理# out: 推理结果 1个 [bs, anchor_num*grid_w*grid_h, xywh+c+20classes] = [1, 19200+4800+1200, 25]# train_out: 训练结果 3个 [bs, anchor_num, grid_w, grid_h, xywh+c+20classes]# 如: [1, 3, 80, 80, 25] [1, 3, 40, 40, 25] [1, 3, 20, 20, 25]out, train_out = model(img, augment=augment)  # inference and training outputst1 += time_synchronized() - t   # 累计前向推理时间 t1

3.7.3、计算验证集损失

        # 6.3、计算验证集损失# compute_loss不为空 说明正在执行train.py 根据传入的compute_loss计算损失值if compute_loss:loss += compute_loss([x.float() for x in train_out], targets)[1][:3]  # lbox, lobj, lcls

3.7.4、Run NMS

        # 6.4、Run NMS# 将真实框target的xywh(因为target是在labelimg中做了归一化的)映射到img(test)尺寸targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device)# save_hybrid: adding the dataset labels to the model predictions before NMS# 是在NMS之前将数据集标签targets添加到模型预测中# 这允许在数据集中自动标记(for autolabelling)其他对象(在pred中混入gt) 并且mAP反映了新的混合标签# targets: [num_target, img_index+class_index+xywh] = [31, 6]# lb: {list: bs} 第一张图片的target[17, 5] 第二张[1, 5] 第三张[7, 5] 第四张[6, 5]lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []t = time_synchronized()out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)t2 += time_synchronized() - t   # 累计NMS时间

3.7.5、统计每张图片的真实框、预测框信息

        # 6.5、统计每张图片的真实框、预测框信息 Statistics per image# 为每张图片做统计，写入预测信息到txt文件，生成json文件字典，统计tp等# out: list{bs} [300, 6] [42, 6] [300, 6] [300, 6] [:, image_index+class+xywh]for si, pred in enumerate(out):# 6.5、统计每张图片的真实框、预测框信息# 获取第si张图片的gt标签信息 包括class, x, y, w, h target[:, 0]为标签属于哪张图片的编号labels = targets[targets[:, 0] == si, 1:]   # [:, class+xywh]nl = len(labels)  # 第si张图片的gt个数# 获取标签类别tcls = labels[:, 0].tolist() if nl else []  # target classpath = Path(paths[si])  # 第si张图片的地址# 统计测试图片数量 +1seen += 1# 如果预测为空，则添加空的信息到stats里if len(pred) == 0:if nl:stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))continue# Predictionsif single_cls:pred[:, 5] = 0predn = pred.clone()# 将预测坐标映射到原图img中scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1])  # native-space pred

3.7.6、保存预测信息到image_name.txt文件

预测信息：cls, xywh, conf

            # 6.6、保存预测信息到txt文件 runs\test\exp7\labels\image_name.txtif save_txt:# gn = [w, h, w, h] 对应图片的宽高 用于后面归一化gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]]  # normalization gain whwhfor *xyxy, conf, cls in predn.tolist():# xyxy -> xywh 并作归一化处理xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywhline = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format# 保存预测类别和坐标值到对应图片image_name.txt文件中with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:f.write(('%g ' * len(line)).rstrip() % line + '\n')# with open(save_dir / 'labels' / ('test' + '.txt'), 'a') as f:# f.write(('%g ' * len(line)).rstrip() % line + '\n')

如下图：

3.7.7、保存预测信息到wandb_logger

            # 6.7、保存预测信息到wandb_logger(类似tensorboard)中if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0:  # Check for test operationif wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:box_data = [{
    "position": {
    "minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},"class_id": int(cls),"box_caption": "%s %.3f" % (names[cls], conf),"scores": {
    "class_score": conf},"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]boxes = {
    "predictions": {
    "box_data": box_data, "class_labels": names}}  # inference-spacewandb_images.append(wandb_logger.wandb.Image(img[si], boxes=boxes, caption=path.name))wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None

3.7.8、将预测信息保存到coco格式的json字典

jdict字典（预测信息）：image_id + category_id + bbox + score

            # 6.8、将预测信息保存到coco格式的json字典(后面存入json文件) Append to pycocotools JSON dictionary# 类似: [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...if save_json:# 获取图片idimage_id = int(path.stem) if path.stem.isnumeric() else path.stem# 获取预测框 并将xyxy转为xywh格式box = xyxy2xywh(predn[:, :4])  # xywh# 之前的的xyxy格式是左上角右下角坐标 xywh是中心的坐标和宽高# 而coco的json格式的框坐标是xywh(左上角坐标 + 宽高)# 所以这行代码是将中心点坐标 -> 左上角坐标 xy center to top-left cornerbox[:, :2] -= box[:, 2:] / 2# image_id: 图片id 即属于哪张图片# category_id: 类别 coco91class()从索引0~79映射到索引0~90# bbox: 预测框坐标# score: 预测得分for p, b in zip(pred.tolist(), box.tolist()):jdict.append({
    'image_id': image_id,'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),'bbox': [round(x, 3) for x in b],'score': round(p[4], 5)})

3.7.9、计算混淆矩阵、计算correct、生成stats

correct(TP)计算过程

for gt中所有类别：

1. 选出pred中属于该类别的所有预测框
2. 选出gt中属于该类别的所有gt框
3. 计算出选出的所有预测框 和 选出的所有gt框 ious
4. 筛选出所有ious > 0.5的预测框 就是TP
5. 如果存在TP 就统计所有TP中不同iou阈值下的TP 同时统计检测到的目标(detected)
6. 重复这个过程 直到检测到的目标个数len(detected) = gt个数

            # 6.9、计算混淆矩阵、计算correct、生成stats# 初始化预测评定 niou为iou阈值的个数 Assign all predictions as incorrect# correct = [pred_obj_num, 10] = [300, 10] 全是Falsecorrect = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)if nl:detected = []  # target indices 用于存放已检测到的目标tcls_tensor = labels[:, 0]  # 当前图片的所有gt的类别tbox = xywh2xyxy(labels[:, 1:5])  # gt boxes 获得xyxy格式的框# 将预测框映射到原图imgscale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1])  # native-space labelsif plots:# 计算混淆矩阵 confusion_matrixconfusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))# Per target class# 对图片中每个类别单独处理for cls in torch.unique(tcls_tensor):# gt中该类别的索引 target indices nonzero: 获取列表中为True的indexti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1)# 预测框中该类别的索引 prediction indicespi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1)# Search for detectionsif pi.shape[0]:# Prediction to target ious# predn[pi, :4]: 属于该类的预测框[144, 4] tbox[ti]: 属于该类的gt框[13, 4]# box_iou: [144, 4] + [13, 4] => [144, 13] 计算属于该类的预测框与属于该类的gt框的iou# .max(1): [144] 选出每个预测框与所有gt box中最大的iou值, i为最大iou值时对应的gt索引ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1)  # best ious, indices# Append detectionsdetected_set = set()  # 这个参数好像没什么用for j in (ious > iouv[0]).nonzero(as_tuple=False):  # j: ious中>0.5的索引 只有iou>=0.5才是TP# 获得检测到的目标d = ti[i[j]]  # detected targetif d.item() not in detected_set:detected_set.add(d.item())  # 没什么用detected.append(d) # 将当前检测到的gt框d添加到detected()# iouv为以0.05为步长 0.5-0.95的序列# 统计所有TP中不同iou阈值下的TP true positive 并在correct中记录下哪个预测框是哪个iou阈值下的TP# correct: [pred_num, 10] = [300, 10] 记录着哪个预测框在哪个iou阈值下是TPcorrect[pi[j]] = ious[j] > iouv  # iou_thres is 1xnif len(detected) == nl:  # 如果检测到的目标值等于gt框的个数 就结束break# 将每张图片的预测结果统计到stats中# Append statistics(correct, conf, pcls, tcls) bs个(correct, conf, pcls, tcls)# correct: [pred_num, 10] bool 当前图片每一个预测框在每一个iou条件下是否是TP# pred[:, 4]: [pred_num, 1] 当前图片每一个预测框的conf# pred[:, 5]: [pred_num, 1] 当前图片每一个预测框的类别# tcls: [gt_num, 1] 当前图片所有gt框的classstats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))

3.7.10、画出前三个batch图片的gt和pred框

        # 6.10、Plot images# 画出前三个batch的图片的ground truth和预测框predictions(两个图)一起保存if plots and batch_i < 3:# ground truthf = save_dir / f'test_batch{
      batch_i}_labels.jpg'# Thread 表示在单独的控制线程中运行的活动 创建一个单线程(子线程)来执行函数 由这个子进程全权负责这个函数# target: 执行的函数 args: 传入的函数参数 daemon: 当主线程结束后, 由他创建的子线程Thread也已经自动结束了# .start(): 启动线程 当thread一启动的时候, 就会运行我们自己定义的这个函数plot_images# 如果在plot_images里面打开断点调试, 可以发现子线程暂停, 但是主线程还是在正常的训练(还是正常的跑)Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start()# predictions 传入plot_images函数之前需要改变pred的格式 target则不需要改f = save_dir / f'test_batch{
      batch_i}_pred.jpg'Thread(target=plot_images, args=(img, output_to_target(out), paths, f, names), daemon=True).start()

test_batch0_labels.jpg

test_batch0_pred.jpg:

3.7.11、计算mAP

    # 6.11、计算mAP# 统计stats中所有图片的统计结果 将stats列表的信息拼接到一起# stats(concat后): list{4} correct, conf, pcls, tcls 统计出的整个数据集的GT# correct [img_sum, 10] 整个数据集所有图片中所有预测框在每一个iou条件下是否是TP [1905, 10]# conf [img_sum] 整个数据集所有图片中所有预测框的conf [1905]# pcls [img_sum] 整个数据集所有图片中所有预测框的类别 [1905]# tcls [gt_sum] 整个数据集所有图片所有gt框的class [929]stats = [np.concatenate(x, 0) for x in zip(*stats)]  # to numpy# stats[0].any(): stats[0]是否全部为False, 是则返回 False, 如果有一个为 True, 则返回 Trueif len(stats) and stats[0].any():# 根据上面的统计预测结果计算p, r, ap, f1, ap_class（ap_per_class函数是计算每个类的mAP等指标的）等指标# p: [nc] 最大平均f1时每个类别的precision# r: [nc] 最大平均f1时每个类别的recall# ap: [71, 10] 数据集每个类别在10个iou阈值下的mAP# f1 [nc] 最大平均f1时每个类别的f1# ap_class: [nc] 返回数据集中所有的类别indexp, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)# ap50: [nc] 所有类别的mAP@0.5 ap: [nc] 所有类别的mAP@0.5:0.95ap50, ap = ap[:, 0], ap.mean(1)# mp: [1] 所有类别的平均precision(最大f1时)# mr: [1] 所有类别的平均recall(最大f1时)# map50: [1] 所有类别的平均mAP@0.5# map: [1] 所有类别的平均mAP@0.5:0.95mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()# nt: [nc] 统计出整个数据集的gt框中数据集各个类别的个数nt = np.bincount(stats[3].astype(np.int64), minlength=nc)  # number of targets per classelse:nt = torch.zeros(1)

3.7.12、print打印各项指标

    # 6.12、print打印各项指标# Print results 数据集图片数量 + 数据集gt框的数量 + 所有类别的平均precision + # 所有类别的平均recall + 所有类别的平均mAP@0.5 + 所有类别的平均mAP@0.5:0.95pf = '%20s' + '%11i' * 2 + '%11.3g' * 4  # print formatprint(pf % ('all', seen, nt.sum(), mp, mr, map50, map))# Print results per class# 细节展示每个类别的各个指标 类别 + 数据集图片数量 + 这个类别的gt框数量 + 这个类别的precision +# 这个类别的recall + 这个类别的mAP@0.5 + 这个类别的mAP@0.5:0.95if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):for i, c in enumerate(ap_class):print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))# Print speeds 打印前向传播耗费的总时间、nms耗费总时间、总时间t = tuple(x / seen * 1E3 for x in (t0, t1, t2))  # speeds per imageif not training:shape = (batch_size, 3, imgsz, imgsz)print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {
      shape}' % t)

输出如图：

3.7.13、画出混淆矩阵并存入wandb_logger中

    # 6.13、画出混淆矩阵并存入wandb_logger中# Plots confusion_matrix + wandb_loggerif plots:confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))if wandb_logger and wandb_logger.wandb:val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]wandb_logger.log({
    "Validation": val_batches})if wandb_images:wandb_logger.log({
    "Bounding Box Debugger/Images": wandb_images})

confusion_matrix.png:

3.7.14、Save JSON

    # 6.14、Save JSON# 采用之前保存的json文件格式预测结果 通过cocoapi评估各个指标# 需要注意的是 测试集的标签也要转为coco的json格式if save_json and len(jdict):w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ''  # weightsanno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json')  # annotations json# 获取预测框的json文件路径并打开pred_json = str(save_dir / f"{
      w}_predictions.json")  # predictions jsonprint('\nEvaluating pycocotools mAP... saving %s...' % pred_json)with open(pred_json, 'w') as f:json.dump(jdict, f)try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynbcheck_requirements(['pycocotools'])from pycocotools.coco import COCOfrom pycocotools.cocoeval import COCOeval# 获取并初始化测试集标签的json文件anno = COCO(anno_json)  # init annotations api# 初始化预测框的文件pred = anno.loadRes(pred_json)  # init predictions api# 创建评估器eval = COCOeval(anno, pred, 'bbox')if is_coco:eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files]  # image IDs to evaluate# 评估eval.evaluate()eval.accumulate()# 展示结果eval.summarize()map, map50 = eval.stats[:2]  # update results (mAP@0.5:0.95, mAP@0.5)except Exception as e:print(f'pycocotools unable to run: {
      e}')

3.7.15、 Return results

    # 6.15、返回测试指标结果 Return resultsmodel.float()  # for trainingif not training:s = f"\n{
      len(list(save_dir.glob('labels/*.txt')))} labels saved to {
      save_dir / 'labels'}" if save_txt else ''print(f"Results saved to {
      save_dir}{
      s}")maps = np.zeros(nc) + map  # [80] 80个平均mAP@0.5:0.95for i, c in enumerate(ap_class):maps[c] = ap[i]  # maps [80] 所有类别的mAP@0.5:0.95# (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()): {tuple:7}# 0: mp [1] 所有类别的平均precision(最大f1时)# 1: mr [1] 所有类别的平均recall(最大f1时)# 2: map50 [1] 所有类别的平均mAP@0.5# 3: map [1] 所有类别的平均mAP@0.5:0.95# 4: val_box_loss [1] 验证集回归损失# 5: val_obj_loss [1] 验证集置信度损失# 6: val_cls_loss [1] 验证集分类损失# maps: [80] 所有类别的mAP@0.5:0.95# t: {tuple: 3} 0: 打印前向传播耗费的总时间 1: nms耗费总时间 2: 总时间return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t

一般会将数据返回到train.py中：

总结

这个脚本的代码还是比较简单的，主要的难点在于第3.7.9节和3.7.11节计算 混淆矩阵+计算correct+计算mAP 上。所以要想全部理解这个脚本，一定要结合metrics.py脚本一起看。

–2021.8.12