yolov5的后处理阶段

yolov5的后处理阶段主要包括以下几个步骤：

1. 对网络输出进行解析，得到预测框的位置、类别和置信度等信息。

# 解析网络输出
pred = model(imgs)  # imgs为输入图像
pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres)

2. 对预测框进行筛选，去除置信度较低的框和重叠度较高的框。

# 对预测框进行筛选
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False):
    ...
    # 筛选置信度较低的框
    scores = pred[:, 4] * max_wh  # 计算置信度
    keep = scores > conf_thres  # 筛选置信度大于阈值的框
    pred = pred[keep]

    # 筛选重叠度较低的框
    if len(pred) > 1:
        # 计算IoU
        iou = bbox_iou(pred[:, :4], pred[:, :4])
        # 获取重叠度较低的框
        iou = torch.triu(iou, diagonal=1)  # 只取上三角部分
        iou_max = iou.max(1)[0]
        keep = (iou_max <= iou_thres)
        pred = pred[keep]
    ...

3. 对预测框进行解码，得到真实框的位置和类别等信息。

# 对预测框进行解码
def xywh2xyxy(x):
    y = x.new(x.shape)
    y[..., 0] = x[..., 0] - x[..., 2] / 2
    y[..., 1] = x[..., 1] - x[..., 3] / 2
    y[..., 2] = x[..., 0] + x[..., 2] / 2
    y[..., 3] = x[..., 1] + x[..., 3] / 2
    return y

def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
    ...
    # 对预测框进行解码
    coords /= ratio
    pad = (img1_shape - ratio_pad) / 2
    coords -= torch.cat([pad, pad], dim=-1)
    coords[:, [0, 2]] = torch.clamp(coords[:, [0, 2]], 0, img0_shape[1])
    coords[:, [1, 3]] = torch.clamp(coords[:, [1, 3]], 0, img0_shape[0])
    return coords

4. 对预测框进行后处理，包括NMS、类别筛选和可视化等。

# 对预测框进行后处理
for i, det in enumerate(pred):  # 遍历每张图像的预测结果
    if det is not None and len(det):
        det[:, :4] = scale_coords(imgs.shape[2:], det[:, :4], img0.shape).round()
        for *xyxy, conf, cls in reversed(det):
            # 绘制预测框
            label = f'{names[int(cls)]} {conf:.2f}'
            plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)