YOLOv5-4.0-test.py 源代码导读(模型测试）

YOLOv5介绍

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

test.py

以下为该文件必须导入的模块，其中部分文件来源于其他项目子文件。utils文件中其他子文件笔者会在后续逐渐更新。 (以下代码均保留原始缩进以方便阅读）

import argparse #解析命令行参数的库
import json     #实现字典列表和JSON字符串之间的相互解析
import os       #与操作系统进行交互的文件库 包含文件路径操作与解析
from pathlib import Path  #Path能够更加方便得对字符串路径进行处理
from threading import Thread #python中处理多线程的库

import numpy as np    #矩阵计算基础库
import torch          #pytorch 深度学习库
import yaml           #yaml是一种表达高级结构的语言 易读 便于指定模型架构及运行配置
from tqdm import tqdm #用于直观显示进度条的一个库 看起来很舒服

from models.experimental import attempt_load #调用models文件夹中的experimental.py文件中的attempt_load函数 目的是加载模型

#以下调用均为utils文件夹中各种已写好的函数
from utils.datasets import create_dataloader
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, box_iou, \
    non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path
from utils.loss import compute_loss
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

测试函数输入变量、加载模型和确定运行设备部分。

#测试函数 输入为测试过程中需要的各种参数
def test(data,
         weights=None,
         batch_size=32,
         imgsz=640,
         conf_thres=0.001,
         iou_thres=0.6,  # 为NMS设置的iou阈值
         save_json=False,
         single_cls=False,
         augment=False,
         verbose=False,
         model=None,
         dataloader=None,
         save_dir=Path(''),  # 存储测试图片的路径
         save_txt=False,  # 自动实现对测试图片的标注
         save_conf=False, #保存置信度
         plots=True,
         log_imgs=0):  # 已记录图片的数量

    # 初始化/加载模型 并设置设备
    training = model is not None #有模型则 training 为True
    if training:  # 调用train.py
        device = next(model.parameters()).device  # 获得记录在模型中的设备 next为迭代器

    else:  
        set_logging() #调用general.py文件中的函数 设置日志 opt对象main中解析传入变量的对象
        device = select_device(opt.device, batch_size=batch_size) #调用torch_utils中select_device来选择执行程序时的设备
        save_txt = opt.save_txt  # 获取保存测试之后的label文件路径 格式为txt

        # 路径
        # 调用genera.py中的increment_path函数来设置保存文件的路径
        save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
        # mkdir创建路径最后一级目录  
        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  

        # 加载模型
        # 加载模型为32位浮点数模型（权重参数） 调用experimental.py文件中的attempt_load函数
        model = attempt_load(weights, map_location=device)
        # 调用general.py中的check_img_size函数来检查图像分辨率能否被32整除  
        imgsz = check_img_size(imgsz, s=model.stride.max())  

确定精度是否减半、加载数据集yaml配置文件、初始化日志部分。

    # 精度减半
    #如果设备类型不是cpu 则将模型由32位浮点数转换为16位浮点数
    half = device.type != 'cpu'  
    if half:
        model.half()

    # 加载配置
    model.eval() #将模型转换为测试模式 固定住dropout层和Batch Normalization层
    is_coco = data.endswith('coco.yaml')  # 判断输入的数据yaml文件是否是coco.yaml文件
    with open(data) as f: #打开data（yaml格式）文件
        data = yaml.load(f, Loader=yaml.FullLoader)  # 获取模型配置的字典格式文件
    check_dataset(data)  # 调用general.py中的check_dataset函数来检查数据文件是否正常
    nc = 1 if single_cls else int(data['nc'])  # 确定检测的类别数目
    iouv = torch.linspace(0.5, 0.95, 10).to(device)  # mAP@0.5:0.95 的iou向量
    niou = iouv.numel() #numel为pytorch预置函数 用来获取张量中的元素个数

    # 日志
    log_imgs, wandb = min(log_imgs, 100), None  # ceil
    try:
        import wandb  # 权重和偏置 wandb为可视化权重和各种指标的库
    except ImportError:
        log_imgs = 0

获取dataloader、初始化模型测试当中用到的相应参数。

    # Dataloader
    if not training:
        img = torch.zeros((1, 3, imgsz, imgsz), device=device)  #创建一张全为0的图片（四维张量）
        # 利用空图片对模型进行测试 （只在运行设备不是cpu时进行）
        _ = model(img.half() if half else img) if device.type != 'cpu' else None  
        path = data['test'] if opt.task == 'test' else data['val']  #如果任务为test 则获得yaml文件中测试的路径
        #调用datasets.py文件中的create_dataloader函数创建dataloader
        dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.5, rect=True)[0]

    seen = 0 #初始化已完成测试的图片数量
    confusion_matrix = ConfusionMatrix(nc=nc) #调用matrics中函数 存储混淆矩阵
    #获取模型训练中存储的类别名字数据
    names = {
   k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
    coco91class = coco80_to_coco91_class()#调用general.py中的函数 来转换coco的类
    #为后续设置基于tqdm的进度条作基础
    s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
    #初始化detection中各个指标的值 t0和t1为时间
    p, r, f1, mp, mr, map50, map, t0, t1