YOLOv5/YOLOv8前后处理

发表于 更新于 分类于
本文字数： 2.7k 阅读时长 ≈ 5 分钟

基本信息

  • YOLOv5: https://github.com/ultralytics/yolov5
  • Commit id: 915bbf294bb74c859f0b41f1c23bc395014ea679
  • Tag: v7.0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
$ git clone  https://github.com/ultralytics/yolov5
$ git checkout -b v7.0 v7.0
$ git log
commit 915bbf294bb74c859f0b41f1c23bc395014ea679 (HEAD -> v7.0, tag: v7.0)
Author: Glenn Jocher <glenn.jocher@ultralytics.com>
Date:   Tue Nov 22 16:23:47 2022 +0100

    YOLOv5 v7.0 release updates (#10245)
    
    * YOLOv5 v7.0 splash image update
    
    * Update tutorial.ipynb
...
...

YOLOv8目前仍在快速迭代中，并且没有发布官方正式版本，所以采用最新提交实现。

  • YOLOv8: https://github.com/ultralytics/ultralytics
  • Commit id: e58db228c2fd9856e7bff54a708bf5acde26fb29
1
2
3
4
5
6
7
8
9
10
11
12
13
14
$ git clone  https://github.com/ultralytics/ultralytics.git
$ git reset --hard e58db228c2fd9856e7bff54a708bf5acde26fb29
$ git log
commit e58db228c2fd9856e7bff54a708bf5acde26fb29 (HEAD -> main)
Author: Glenn Jocher <glenn.jocher@ultralytics.com>
Date:   Thu Oct 26 20:32:51 2023 +0200

    `ultralytics 8.0.202` sort Triton model outputs (#5945)
    
    Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
    Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
    Co-authored-by: Mike Tune <mtuneoff@gmail.com>
...
...

YOLOv5

前处理

读取图片

在v5工程的批量数据读取流程中，输出的targets是大小为[N, 6]的二维数组。每个数组项（item）都包含以下六个元素：

1
[img_id, cls_id, x_c, y_c, box_w, box_h]

其中，img_id表示图像索引，它是在dataloder通过静态方法collate_fn设置的。

1
2
3
4
5
6
7
8
9
10
11
12
# dataloders.py 844
        @staticmethod  
        def collate_fn(batch):  
            # 将batch中的元素进行解包，并分别赋值给im, label, path, shapes  
            im, label, path, shapes = zip(*batch)  
          
            # 对label列表中的每一项（每个标签数组），更新其第一个元素为当前图像的索引  
            for i, lb in enumerate(label):  
                lb[:, 0] = i  # 为build_targets()方法添加目标图像的索引  
          
            # 将图像数据堆叠为一个tensor，标签数据拼接为一个tensor，并返回图像、标签、路径和形状信息  
            return torch.stack(im, 0), torch.cat(label, 0), path, shapes

数据填充

对于<填充+缩放>操作，YOLOv5提供了3种填充方式：

  • 原图

  • 正常模式（默认设置）：auto=False, scaleFill=False（等比缩放+较短边填充）

  • auto模式：auto=True, scaleFill=False（等比缩放，较短边填充大小为dh=dh%stride）
    • 注意：仅在使用pytorch进行训练和推理时，设置auto=True

  • scaleFill模式：auto=False, scaleFill=True（直接缩放，没有填充操作）

数据归一化

1
2
3
4
# train.py 285
            imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0
# detect.py 119
            im /= 255  # 0 - 255 to 0.0 - 1.0

后处理

在后处理阶段，主要执行置信度过滤和IOU过滤，其中

  1. 置信度过滤实现了两次，首先是目标置信度过滤，然后计算置信度，执行置信度过滤；
  2. IOU过滤需要按照置信度大小降序后依次遍历，保证结果列表中不存在目标框IOU大于阈值的预测结果。

为了能够批量执行NMS，不同类别的坐标框会增加不一样的距离，来保证NMS过程中不同类别的坐标框不会重叠。

1
2
3
4
# Batched NMS
c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
i = nms(boxes, scores, iou_thres)  # NMS

YOLOv5还提供了一些可选实现，包括

  1. 仅保留指定类别预测框；
  2. 不分类别执行NMS；
  3. 最大预测框数目约束。

YOLOv8

前处理

YOLOv8的前处理实现逻辑跟YOLOv5一样，只是在实现步骤以及具体实现流程上有差别。比如

  • YOLOv8将填充+缩放定义为类LetterBox实现，而YOLOv5使用函数letterbox实现；
  • YOLOv5先执行通道维度转换（HWC2CHW），再执行颜色空间转换（BGR2RGB），而YOLOv8相反。

后处理

  • 差异点一：YOLOv8模型采用了Anchor-Free模式，其输出的特征向量中并没有目标置信度，所以不需要执行目标置信度过滤；
  • 差异点二：YOLOv8后处理指定了最大参与NMS的预测框数目（max_nms=30000），在执行完置信度过滤，初步生成输出格式(xyxy, conf, j)之后，先进行max_nms截断，再执行NMS，最后是最大预测框数目（max_det=300）截断后输出。