yolov5检测代码解析

yolov5源代码：https://github.com/ultralytics/yolov5
设置目标检测的配置参数：

parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    opt = parser.parse_args()
    print(opt)

执行检测程序前，查看requirements.txt里的相关安装包是否安装在执行检测的python环境下：

check_requirements(exclude=('pycocotools', 'thop'))

def check_requirements(file='requirements.txt', exclude=()):
    # Check installed dependencies meet requirements
    import pkg_resources as pkg
    prefix = colorstr('red', 'bold', 'requirements:')
    file = Path(file)
    if not file.exists():
        print(f"{prefix} {file.resolve()} not found, check failed.")
        return

    n = 0  # number of packages updates
    requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude]
    for r in requirements:
        try:
            pkg.require(r)
        except Exception as e:  # DistributionNotFound or VersionConflict if requirements not met
            n += 1
            print(f"{prefix} {e.req} not found and is required by YOLOv5, attempting auto-update...")
            print(subprocess.check_output(f"pip install '{e.req}'", shell=True).decode())

    if n:  # if packages updated
        s = f"{prefix} {n} package{'s' * (n > 1)} updated per {file.resolve()}\n" \
            f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
        print(emojis(s))  # emoji-safe

执行检测部分：

    with torch.no_grad():
        if opt.update:  # update all models (to fix SourceChangeWarning)
            for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
                detect()
                strip_optimizer(opt.weights)
        else:
            detect()

接下来主要看detect()内部的执行。
第一行是相关参数的加载，参数加载的内容为输入的配置参数。
第二行为是否保存图片的判断（bool类型），默认为保存图片，不保存图片可以在执行程序中输入 --nosave，或者你检测的对象是一个网络视频流的地址，其中视频流地址保存在一个文件类型为txt的文档里。
第三个是摄像头检测的一个判断（bool类型）， 是你在–source后的输入，当检测对象是本地的一个摄像头或者是一个txt文档内有网络视频源地址，或者是你直接输入网络视频源的地址以(‘rtsp://’, ‘rtmp://’, ‘http://’)这三个开头的网络地址，以上三种情况（不是指网络地址的三种情况）判定你的检测对象是摄像头检测。

def detect(save_img=False):
    source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
    save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
    webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
        ('rtsp://', 'rtmp://', 'http://'))

保存检测结果地址的一个写入程序：
第一行为保存地址返回（str），其中opt.project，默认为’runs/detect’， opt.name默认为’exp’，opt.exist_ok默认为False，opt.exist_ok判断是否将检测结果保存在你默认或者输入的opt.project/opt.name文件内，执行检测程序时如果不改变opt.project和opt.name，只输入–exist_ok，则保存结果会在’runs/detect/exp文件夹内。如果不输入默认的opt.project，opt.name和opt.exist_ok ，increment_path()这个程序是通过通过正则匹配的方式在你的runs/detect文件内查找含有exp的文件，并匹配出最大的数字，进行+1后，返回’runs/detect/exp’ + '(maxnum+1)'的一个字符串
第二行为第一行返回保存位置的一个文件夹创建

    # Directories
    save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run
    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

def increment_path(path, exist_ok=True, sep=''):
    # Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc.
    path = Path(path)  # os-agnostic
    if (path.exists() and exist_ok) or (not path.exists()):
        return str(path)
    else:
        dirs = glob.glob(f"{path}{sep}*")  # similar paths
        matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
        i = [int(m.groups()[0]) for m in matches if m]  # indices
        n = max(i) + 1 if i else 2  # increment number
        return f"{path}{sep}{n}"  # update path

第一行打印日志信息
第二行选择程序在哪个硬件执行（CPU或者GPU），select_device()，如果在opt.device选择输入cpu，则在cpu上进行检测程序的执行，如果不输入，则通过torch.cuda.is_available()来判断是否使用GPU运行程序
第三行通过降低少许检测精度提高检测速度的判断（bool）

    # Initialize
    set_logging()
    device = select_device(opt.device)
    half = device.type != 'cpu'  # half precision only supported on CUDA

def select_device(device='', batch_size=None):
    # device = 'cpu' or '0' or '0,1,2,3'
    s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} '  # string
    cpu = device.lower() == 'cpu'
    if cpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # force torch.cuda.is_available() = False
    elif device:  # non-cpu device requested
        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
        assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested'  # check availability

    cuda = not cpu and torch.cuda.is_available()
    if cuda:
        n = torch.cuda.device_count()
        if n > 1 and batch_size:  # check that batch_size is compatible with device_count
            assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
        space = ' ' * len(s)
        for i, d in enumerate(device.split(',') if device else range(n)):
            p = torch.cuda.get_device_properties(i)
            s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n"  # bytes to MB
    else:
        s += 'CPU\n'

    logger.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s)  # emoji-safe
    return torch.device('cuda:0' if cuda else 'cpu')

第一行为模型的载入，其中attempt_load()是对权重文件的一个载入，主要是通过torch.load()这个函数，值得注意的是，这个可以是多个模型的混和载入（需要注意每个模型的输入和输入大小），也可以是对一个完整检测模型的载入。
第二行是检测模型时yolov5也用了FPN里面的一个思路，将卷积神经网络中最后三层通过上采样的方式进行了一个堆叠的，要求图片的分辨率为32的一个整数倍。
第三行为输入判断你输入的opt.img_size是否是32的倍数，如果不是，则改变你输入的img_size为向上取32的倍数的分辨率。

    # Load model
    model = attempt_load(weights, map_location=device)  # load FP32 model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check img_size
    if half:
        model.half()  # to FP16

def attempt_load(weights, map_location=None):
    # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
    model = Ensemble()
    for w in weights if isinstance(weights, list) else [weights]:
        attempt_download(w)
        ckpt = torch.load(w, map_location=map_location)  # load
        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model

    # Compatibility updates
    for m in model.modules():
        if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
            m.inplace = True  # pytorch 1.7.0 compatibility
        elif type(m) is Conv:
            m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility

    if len(model) == 1:
        return model[-1]  # return model
    else:
        print('Ensemble created with %s\n' % weights)
        for k in ['names', 'stride']:
            setattr(model, k, getattr(model[-1], k))
        return model  # return ensemble

def check_img_size(img_size, s=32):
    # Verify img_size is a multiple of stride s
    new_size = make_divisible(img_size, int(s))  # ceil gs-multiple
    if new_size != img_size:
        print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
    return new_size

是否进行第二步检测分类，通过resnet101对通过yolov5识别检测后的结果在进行一次预测。

# Second-stage classifier
    classify = False
    if classify:
        modelc = load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()

数据集的载入：
第一行是声明视频保存和vid_writer两个变量。
第二行，if里是如果检测目标是摄像头或者网络视频源，check_imshow()检测运行环境是否是在docker下（docker环境下不能打开界面显示，会报错），然后是视频流或者摄像头的一个载入。
如果不是视频流或者摄像头，载入检测对象为图片、视频或者文件夹内的文件，通过迭代的方式一个一个的进行检测。

    # Set Dataloader
    vid_path, vid_writer = None, None
    if webcam:
        view_img = check_imshow()
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
    else:
        dataset = LoadImages(source, img_size=imgsz, stride=stride)

载入模型内部的class names，以及当前检测每一个name对应的矩形框的颜色

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

开始执行检测推理部分的代码：
第一个if内的代码是运行一次模型推理，推理为一张全为数值0的图片，图片分辨率为输入的imgsz* imgsz（opencv里的话是全黑的图片）。主要可能是为了检测模型检测是否会报错和让GPU开始预热，哈哈哈。
t0是为了计算整个推理过程的时间的开始计时，即完成文件夹内图像的推理时间。
dataset里的内容可以参考下LoadImages类函数里的__next__()，并通过__iter__()，进行文件夹内部文件的迭代更新。
for里的内容为先将图片内的数值载入到的device（GPU or CPU），如果是在GPU上进行推理（前面的half）将像数值的uint8类型转为fp16或者fp32，然后将像数值的变化范围固定到0.0-1.0之间（对深度学习推理计算更友好吧），最后判断图片的shape是不是3，因为在tensor环境下使用的是ndimension()，如果不是则增加一层。
path为检测对象的地址，img为将检测的原始图片resize到目标尺寸（不失真放缩，通过填充RBG为114, 114, 114的颜色条）

    # Run inference
    if device.type != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    t0 = time.time()
    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

t1是计算单个文件的推理时间的开始计时，在GPU上通过torch.cuda.synchronize()使推理时间计算的更加精确。
pred是模型推理输出的全部结果，以输入尺寸图片为640640为例，coco80类检测，输出的pred的shape为2520085，（20203）大分辨率下的2020个滑块同时乘以大分辨率对应的3个anchor；（40403）中分辨率下的4040个滑块同时乘以中分辨率对应的3个anchor；（80803）小分辨率下的80*80个滑块同时乘以小分辨率对应的3个anchor；三组数相加为25200，85其中[0:4]为anchor画出来的矩形框，[4]为ojbconf即anchor中含有80类中目标检测的置信度，后面的[5:]为80类别对应的每一类的置信度。

 		# Inference
        t1 = time_synchronized()
        pred = model(img, augment=opt.augment)[0]
def time_synchronized():
    # pytorch-accurate time
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    return time.time()

极大值抑制，non_max_suppression()函数太长了，就不在这里贴出来了，有兴趣的可以在源码里看看，在这里我大致说下极大值抑制的流程吧。
第一步是判断pred列表里ojbconf也就是pred中其中一个anchor中[4]小于opt.conf_thres的值，返回一个True和Flase的列表，长度为25200。
当遇到判断列表中的第一个True，先计算类别置信度，即obj_conf*80类别算法得出对应的置信度，再将前面4个数即anchor对应的矩形框转换为xyxy的矩形框（box）的像数点坐标值。然后取出得到置信度（conf）中最大的值和所在的位置（classID），再将box的xyxy、conf以及classID拼凑起来为一个列表。最后通过NMS极大值抑制，通过对计算每个为True的box的交并比，输入满足要求的输入阈值的一个list，list的shape[1]为6,[x,y,x,y,conf,classID]。

        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        t2 = time_synchronized()

是否对预测结果进行二次预测分类

        # Apply Classifier
        if classify:
            pred = apply_classifier(pred, modelc, img, im0s)

预测结果的处理：
首先是判断预测的对象是是为视频源还是其他的对象，设声明结果保存的地址（图片结果以及标签结果），s为每张resize到目标尺寸图像的shape，即喂入网络的图片尺寸（str），在tensor上获得原始检测画面的size。

# Process detections
        for i, det in enumerate(pred):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
            else:
                p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path
            save_path = str(save_dir / p.name)  # img.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

如果det中有相应的检测目标，将det[:, 0:4]，即在img（原图（im0s）resize到targetSize后的图像）上xyxy对于的像数值坐标还原到原图（im0s）的像数值坐标，为了后续方便在原图上画出目标物的矩形框。
然后是确定det[:, -1]中的所存在同类的classID一共有多少个，并有几个类别，是在后续输出检测完成后打印检测结果的相关信息。
检测结果写入，如果在程序执行时输入了 --save-txt,则将类别 xywh 置信度 写入保存在save_dir / ‘labels’/图片名中，如果去掉置信度和前后的空格就是通过labelimg标注后的yolo标签。
执行程序的时候，没有设置 --nosave或设置了–view-img，则在im0s原始图像上画出预测结果的矩形框，为了后续的保存图像、视频结果或者网络视频源。
在命令框内打印检测总体结果

			if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or view_img:  # Add bbox to image
                        label = f'{names[int(cls)]} {conf:.2f}'
                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

            # Print time (inference + NMS)
            print(f'{s}Done. ({t2 - t1:.3f}s)')

检测时需要查看图片或者视频结果，通过imshow显示，显示检测后的图像或视频结果。

            # Stream results
            if view_img:
                cv2.imshow(str(p), im0)
                cv2.waitKey(1)  # 1 millisecond

当检测对象是图片是，保存图片到save_path，为save_dir / 完整的图片名（有尾缀）。
当检测对象是视频或者视频源时，因为对视频检测是对视频中的每一帧的图像进行检测，检测视频显示是通过将视频的每一帧拼接成视频显示出来的，当视频检测结果保存的save_path在同一个视频的时候不会发生变化，当发生变化时，cv2.VideoWriter里的save_path也需要相应的改变。
如果检测对象是本地视频则通过cv2获取fps，w，h，检测对象是网络视频默认fps=30。
然后通过cv2.VideoWriter(args).write（im0）写入每一帧检测结果，保存视频。

            if save_img:
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                else:  # 'video' or 'stream'
                    if vid_path != save_path:  # new video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # release previous video writer
                        if vid_cap:  # video
                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        else:  # stream
                            fps, w, h = 30, im0.shape[1], im0.shape[0]
                            save_path += '.mp4'
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                    vid_writer.write(im0)

打印最总的检测结果。

    if save_txt or save_img:
        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}")

    print(f'Done. ({time.time() - t0:.3f}s)')

欢迎讨论，并指出纠正上面的错误。