2021SC@SDUSC
目录
导入第三方库
设置opt参数
main函数
run函数
import argparse
import json
import os
import sys
from pathlib import Path
from threading import Thread
import numpy as np
import torch
from tqdm import tqdm
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
from models.experimental import attempt_load
from utils.datasets import create_dataloader
from utils.general import coco80_to_coco91_class, check_dataset, check_img_size, check_requirements, \
check_suffix, check_yaml, box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, \
increment_path, colorstr, print_args
from utils.metrics import ap_per_class, ConfusionMatrix
from utils.plots import output_to_target, plot_images, plot_val_study
from utils.torch_utils import select_device, time_sync
from utils.callbacks import Callbacks
def parse_opt():
parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='dataset.yaml path')
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--batch-size', type=int, default=32, 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', 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', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')
parser.add_argument('--project', default='runs/val', 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')
opt.save_txt |= opt.save_hybrid
opt.data = check_yaml(opt.data) # check YAML
print_args(FILE.stem, opt)
return opt
data: 数据集配置文件地址 包含数据集的路径、类别个数、类名、下载地址等信息
weights: 模型的权重文件地址 weights/yolov5s.pt
batch_size: 前向传播的批次大小 默认32
imgsz: 输入网络的图片分辨率 默认640
conf-thres: object置信度阈值 默认0.25
iou-thres: 进行NMS时IOU的阈值 默认0.6
task: 设置测试的类型 有train, val, test, speed or study几种 默认val
device: 测试的设备
single-cls: 数据集是否只用一个类别 默认False
augment: 测试是否使用TTA Test Time Augment 默认False
verbose: 是否打印出每个类别的mAP 默认False
save-txt: traditional auto-labelling
save-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 commands
save-json: 是否按照coco的json格式保存预测框,并且使用cocoapi做评估(需要同样coco的json格式的标签) 默认
False
project: 测试保存的源文件 默认runs/test
name: 测试保存的文件地址 默认exp 保存在runs/test/exp下
exist-ok: 是否存在当前文件 默认False 一般是 no exist-ok 连用 所以一般都要重新创建文件夹
half: 是否使用半精度推理 默认False
def main(opt):
# 初始化日志 设置logging的等级level、输出格式和内容format等基本信息
set_logging()
# 检测requirements文件中需要的包是否安装好了
check_requirements(exclude=('tensorboard', 'thop'))
# 如果task in ['train', 'val', 'test']就正常测试 训练集/验证集/测试集
if opt.task in ('train', 'val', 'test'): # run normally
run(**vars(opt))
# 如果opt.task == 'speed' 就测试yolov5系列和yolov3-spp各个模型的速度评估
elif opt.task == 'speed': # speed benchmarks
for 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.pt
x = 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 to
y = [] # y axis
for i in x: # img-size
print(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 times
np.savetxt(f, y, fmt='%10.4g') # save
os.system('zip -r study.zip study_*.txt')
plot_val_study(x=x) # plot
可以看到这个模块根据opt.task可以分为三个分支,我们主要的分支还是在 opt.task in (‘train’, ‘val’, ‘test’)。而其他的两个分支,大家大概看看在干什么就可以了,没什么用。一般我们都是直接进入第一个分支,执行run函数。
def run(data,
weights=None, # model.pt path(s)
batch_size=32, # batch size
imgsz=640, # inference size (pixels)
conf_thres=0.001, # confidence threshold
iou_thres=0.6, # NMS IoU threshold
task='val', # train, val, test, speed or study
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
single_cls=False, # treat as single-class dataset
augment=False, # augmented inference
verbose=False, # verbose output
save_txt=False, # save results to *.txt
save_hybrid=False, # save label+prediction hybrid results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_json=False, # save a COCO-JSON results file
project='runs/val', # save to project/name
name='exp', # save to project/name
exist_ok=False, # existing project/name ok, do not increment
half=True, # use FP16 half-precision inference
model=None,
dataloader=None,
save_dir=Path(''),
plots=True,
callbacks=Callbacks(),
compute_loss=None,
):
data: 数据集配置文件地址 包含数据集的路径、类别个数、类名、下载地址等信息 train.py时传入data_dict
weights: 模型的权重文件地址 运行train.py=None 运行test.py=默认weights/yolov5s.pt
batch_size: 前向传播的批次大小 运行test.py传入默认32 运行train.py则传入batch_size // WORLD_SIZE * 2
imgsz: 输入网络的图片分辨率 运行test.py传入默认640 运行train.py则传入imgsz_test
conf_thres: object置信度阈值 默认0.25
iou_thres: 进行NMS时IOU的阈值 默认0.6
task: 设置测试的类型 有train, val, test, speed or study几种 默认val
device: 测试的设备
single_cls: 数据集是否只用一个类别 运行test.py传入默认False 运行train.py则传入single_cls
augment: 测试是否使用TTA Test Time Augment 默认False
verbose: 是否打印出每个类别的mAP 运行test.py传入默认Fasle 运行train.py则传入nc < 50 and final_epoch
save_txt: 是否以txt文件的形式保存模型预测框的坐标 默认True
save_hybrid: 是否save label+prediction hybrid results to *.txt 默认False
save_conf: 是否保存预测每个目标的置信度到预测tx文件中 默认True
save_json: 是否按照coco的json格式保存预测框,并且使用cocoapi做评估(需要同样coco的json格式的标签)
project: 测试保存的源文件 默认runs/test
name: 测试保存的文件地址 默认exp 保存在runs/test/exp下
exist_ok: 是否存在当前文件 默认False 一般是 no exist-ok 连用 所以一般都要重新创建文件夹
half: 是否使用半精度推理 FP16 half-precision inference 默认False
model: 模型 如果执行test.py就为None 如果执行train.py就会传入ema.ema(ema模型)
dataloader: 数据加载器 如果执行test.py就为None 如果执行train.py就会传入testloader
save_dir: 文件保存路径 如果执行test.py就为‘’ 如果执行train.py就会传入save_dir(runs/train/expn)
plots: 是否可视化 运行test.py传入默认True 运行train.py则传入plots and final_epoch
wandb_logger: 网页可视化 类似于tensorboard 运行test.py传入默认None 运行train.py则传入wandb_logger(train)
compute_loss: 损失函数 运行test.py传入默认None 运行train.py则传入compute_loss(train)