导航

• data
  • coco128.yaml
  • hyp.scratch-low.yaml
• train.py
• val.py
• detect.py

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

import argparse  # 解析命令行参数模块
import json    # 实现字典列表和JSON字符串之间的相互解析
import os   # 与操作系统进行交互的模块 包含文件路径操作和解析
import sys   # sys系统模块 包含了与Python解释器和它的环境有关的函数
from pathlib import Path # Path将str转换为Path对象 使字符串路径易于操作的模块

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
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend
from utils.callbacks import Callbacks
from utils.dataloaders import create_dataloader
from utils.general import (LOGGER, TQDM_BAR_FORMAT, Profile, check_dataset, check_img_size, check_requirements,
                           check_yaml, coco80_to_coco91_class, colorstr, increment_path, non_max_suppression,
                           print_args, scale_boxes, xywh2xyxy, xyxy2xywh)
from utils.metrics import ConfusionMatrix, ap_per_class, box_iou
from utils.plots import output_to_target, plot_images, plot_val_study
from utils.torch_utils import select_device, smart_inference_mode

1、设置opt参数

    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') # 数据集配置文件地址 包含数据集的路径、类别个数、类名、下载地址等信息
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)') #  模型的权重文件地址yolov5s.pt
    parser.add_argument('--batch-size', type=int, default=32, help='batch size') # 前向传播的批次大小 默认32
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)') # 输入网络的图片分辨率 默认640
    parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold') # object置信度阈值 默认0.001
    parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')  # 进行NMS时IOU的阈值 默认0.6
    parser.add_argument('--max-det', type=int, default=300, help='maximum detections per image')
    parser.add_argument('--task', default='val', help='train, val, test, speed or study')  # 设置测试的类型 有train, val, test, speed or study几种 默认val
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') # 测试的设备
    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
    parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') #  数据集是否只用一个类别 默认False
    parser.add_argument('--augment', action='store_true', help='augmented inference') # 测试是否使用TTA Test Time Augment 默认False
    parser.add_argument('--verbose', action='store_true', help='report mAP by class')  # 是否打印出每个类别的mAP 默认False
    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') #  是否按照coco的json格式保存预测框，并且使用cocoapi做评估（需要同样coco的json格式的标签） 默认False
    parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name') # 测试保存的源文件 默认runs/val
    parser.add_argument('--name', default='exp', help='save to project/name') # 测试保存的文件地址 默认exp  保存在runs/val/exp下
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') # 是否存在当前文件 默认False 一般是 no exist-ok 连用  所以一般都要重新创建文件夹
    parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') #  是否使用半精度推理 默认False
    parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
    opt = parser.parse_args()
    opt.data = check_yaml(opt.data)  # check YAML
    opt.save_json |= opt.data.endswith('coco.yaml')
    opt.save_txt |= opt.save_hybrid
    print_args(vars(opt))
    return opt