yolov5s后处理 python实现

yolov5s的后处理包括NMS（非极大值抑制）和解码。下面是一个使用Python实现yolov5s后处理的例子：

import torch
from torchvision.ops import nms
from utils.general import non_max_suppression

def yolov5s_postprocess(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False):
    # 获取预测结果
    prediction = non_max_suppression(prediction, conf_thres, iou_thres, classes=classes, agnostic=agnostic)

    # 遍历每个预测结果
    for i, det in enumerate(prediction):
        # 如果没有预测结果，跳过
        if det is None:
            continue

        # 对预测结果进行NMS
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
        det = nms(det[:, :4], det[:, 4], iou_thres)

        # 遍历每个预测框
        for *xyxy, conf, cls in reversed(det):
            # 打印预测结果
            print(f'Object {i}: {names[int(cls)]} ({conf:.2f})')

            # 绘制预测框
            label = f'{names[int(cls)]} {conf:.2f}'
            plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

    return prediction

其中，`prediction`是yolov5s的预测结果，`conf_thres`是置信度阈值，`iou_thres`是IoU阈值，`classes`是要检测的类别，`agnostic`表示是否使用类别不可知的NMS。这个函数会返回处理后的预测结果。