YOLOV5测试代码test.py注释与解析

YOLOv5代码注释版更新啦，注释的是最近的2021.07.14的版本，且注释更全
github: https://github.com/Laughing-q/yolov5_annotations

本文主要对ultralytics\yolov5的测试代码test.py的解析，由于yolov5还在开发当中，平常多多少少都会修复一些bug或者有一些代码和功能的更新，但基本上不会有很大的改动，故以下注释与解析都是适用的；当然如果有大改动，笔者也会更新注释。

yolov5其他代码解析

测试参数以及main函数解析

if __name__ == '__main__':
    """
    opt参数详解
    weights:测试的模型权重文件
    data:数据集配置文件，数据集路径，类名等
    batch-size:前向传播时的批次, 默认32
    img-size:输入图片分辨率大小, 默认640
    conf-thres:筛选框的时候的置信度阈值, 默认0.001
    iou-thres:进行NMS的时候的IOU阈值, 默认0.65
    save-json:是否按照coco的json格式保存预测框，并且使用cocoapi做评估(需要同样coco的json格式的标签), 默认False
    task:设置测试形式, 默认val, 具体可看下面代码解析注释
    device:测试的设备，cpu；0(表示一个gpu设备cuda:0)；0,1,2,3(多个gpu设备)
    single-cls:数据集是否只有一个类别，默认False
    augment:测试时是否使用TTA(test time augmentation), 默认False
    merge:在进行NMS时，是否通过合并方式获得预测框, 默认False
    verbose:是否打印出每个类别的mAP, 默认False
    save-txt:是否以txt文件的形式保存模型预测的框坐标, 默认False
    """
    parser = argparse.ArgumentParser(prog='test.py')
    parser.add_argument('--weights', nargs='+', type=str, default='runs/exp0/weights/last.pt', help='model.pt path(s)')
    parser.add_argument('--data', type=str, default='data/mask.yaml', help='*.data path')
    parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
    parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.60, help='IOU threshold for NMS')
    parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
    parser.add_argument('--task', default='val', help="'val', 'test', '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('--merge', action='store_true', help='use Merge NMS')
    parser.add_argument('--verbose', action='store_true', help='report mAP by class')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    opt = parser.parse_args()
    # 设置参数save_json
    opt.save_json = opt.save_json or opt.data.endswith('coco.yaml')
    # check_file检查文件是否存在
    opt.data = check_file(opt.data)  # check file
    print(opt)

    # task = 'val', 'test', 'study'
    # task = ['val', 'test']时就正常测试验证集、测试集
    if opt.task in ['val', 'test']:  # (default) run normally
        test(opt.data,
             opt.weights,
             opt.batch_size,
             opt.img_size,
             opt.conf_thres,
             opt.iou_thres,
             opt.save_json,
             opt.single_cls,
             opt.augment,
             opt.verbose)

    # task == 'study'时，就评估yolov5系列和yolov3-spp各个模型在各个尺度下的指标并可视化
    elif opt.task == 'study':  # run over a range of settings and save/plot
        for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt']:
            f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem)  # filename to save to
            x = list(range(352, 832, 64))  # x axis
            y = []  # y axis
            for i in x:  # img-size
                print('\nRunning %s point %s...' % (f, i))
                r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json)
                y.append(r + t)  # results and times
            np.savetxt(f, y, fmt='%10.4g')  # save
        os.system('zip -r study.zip study_*.txt')
        plot_study_txt(f, x)  # plot

test函数解析

import argparse
import json

from models.experimental import *