YOLOv5推理代码解析

代码如下

import cv2
import numpy as np
import onnxruntime as ort
import time
import random

# 画一个检测框
def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    description: 在图像上绘制一个矩形框。
    param:
        x: 框的坐标 [x1, y1, x2, y2]
        img: 输入图像
        color: 矩形框的颜色，默认为随机颜色
        label: 框内显示的标签
        line_thickness: 矩形框的线条宽度
    return: 无返回值，直接在图像上绘制
    """
    tl = (
            line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness，计算线条或字体的粗细
    color = color or [random.randint(0, 255) for _ in range(3)]  # 如果没有提供颜色，随机生成颜色
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))  # 左上角和右下角的坐标
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)  # 绘制矩形框
    if label:  # 如果提供了标签，则绘制标签
        tf = max(tl - 1, 1)  # 字体的粗细
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]  # 获取标签的大小
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3  # 计算标签背景框的位置
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # 绘制标签背景框
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
        )  # 绘制标签文本

# 生成网格坐标
def _make_grid(nx, ny):
    """
    description: 生成网格坐标，用于解码预测框位置。
    param:
        nx, ny: 网格的行数和列数
    return: 返回网格坐标
    """
    xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))  # 生成网格坐标
    return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32)  # 转换为需要的格式

# 输出解码
def cal_outputs(outs, nl, na, model_w, model_h, anchor_grid, stride):
    """
    description: 对模型输出的坐标进行解码，转换为图像坐标。
    param:
        outs: 模型输出的框的偏移量
        nl: 输出层数量
        na: 每层的anchor数目
        model_w, model_h: 模型输入图像的尺寸
        anchor_grid: anchor的尺寸
        stride: 每个输出层的缩放步长
    return: 解码后的输出
    """
    row_ind = 0
    grid = [np.zeros(1)] * nl  # 每个层对应一个网格
    for i in range(nl):
        h, w = int(model_w / stride[i]), int(model_h / stride[i])  # 计算该层特征图的高和宽
        length = int(na * h * w)  # 当前层的总框数
        if grid[i].shape[2:4] != (h, w):  # 如果网格的大小不匹配，则重新生成网格
            grid[i] = _make_grid(w, h)

        # 解码每个框的中心坐标和宽高
        outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile(
            grid[i], (na, 1))) * int(stride[i])
        outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat(
            anchor_grid[i], h * w, axis=0)  # 计算宽高
        row_ind += length
    return outs

# 后处理，计算检测框
def post_process_opencv(outputs, model_h, model_w, img_h, img_w, thred_nms, thred_cond):
    """
    description: 对模型输出的框进行后处理，得到最终的检测框。
    param:
        outputs: 模型输出的框
        model_h, model_w: 模型输入的高度和宽度
        img_h, img_w: 原图的高度和宽度
        thred_nms: 非极大值抑制的阈值
        thred_cond: 置信度阈值
    return: 返回处理后的框、置信度和类别
    """
    conf = outputs[:, 4].tolist()  # 获取每个框的置信度
    c_x = outputs[:, 0] / model_w * img_w  # 计算中心点x坐标
    c_y = outputs[:, 1] / model_h * img_h  # 计算中心点y坐标
    w = outputs[:, 2] / model_w * img_w  # 计算框的宽度
    h = outputs[:, 3] / model_h * img_h  # 计算框的高度
    p_cls = outputs[:, 5:]  # 获取分类得分
    if len(p_cls.shape) == 1:  # 如果分类结果只有一维，增加一维
        p_cls = np.expand_dims(p_cls, 1)
    cls_id = np.argmax(p_cls, axis=1)  # 获取类别编号

    # 计算框的四个角坐标
    p_x1 = np.expand_dims(c_x - w / 2, -1)
    p_y1 = np.expand_dims(c_y - h / 2, -1)
    p_x2 = np.expand_dims(c_x + w / 2, -1)
    p_y2 = np.expand_dims(c_y + h / 2, -1)
    areas = np.concatenate((p_x1, p_y1, p_x2, p_y2), axis=-1)  # 合并成框的坐标

    areas = areas.tolist()  # 转为列表形式
    ids = cv2.dnn.NMSBoxes(areas, conf, thred_cond, thred_nms)  # 非极大值抑制

    if len(ids) > 0:  # 如果有框被保留
        return np.array(areas)[ids], np.array(conf)[ids], cls_id[ids]
    else:
        return [], [], []

# 图像推理
def infer_img(img0, net, model_h, model_w, nl, na, stride, anchor_grid, thred_nms=0.4, thred_cond=0.5):
    """
    description: 对输入图像进行推理，输出检测框。
    param:
        img0: 原始图像
        net: 加载的ONNX模型
        model_h, model_w: 模型的输入尺寸
        nl: 输出层数量
        na: 每层的anchor数量
        stride: 每层的缩放步长
        anchor_grid: 每层的anchor尺寸
        thred_nms: 非极大值抑制阈值
        thred_cond: 置信度阈值
    return: 检测框、置信度和类别
    """
    # 图像预处理
    img = cv2.resize(img0, [model_w, model_h], interpolation=cv2.INTER_AREA)  # 将图像调整为模型输入大小
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 转换为RGB格式
    img = img.astype(np.float32) / 255.0  # 归一化
    blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)  # 将图像转为模型输入格式

    # 模型推理
    outs = net.run(None, {net.get_inputs()[0].name: blob})[0].squeeze(axis=0)  # 推理并去掉batch维度

    # 输出坐标矫正
    outs = cal_outputs(outs, nl, na, model_w, model_h, anchor_grid, stride)

    # 检测框计算
    img_h, img_w, _ = np.shape(img0)  # 获取原图的尺寸
    boxes, confs, ids = post_process_opencv(outs, model_h, model_w, img_h, img_w, thred_nms, thred_cond)

    return boxes, confs, ids

if __name__ == "__main__":
    # 加载ONNX模型
    model_pb_path = "a.onnx"  # 模型文件路径
    so = ort.SessionOptions()
    net = ort.InferenceSession(model_pb_path, so)

    # 类别字典
    dic_labels = {
        0: 'jn', 1: 'pill_bag', 2: 'pill_ban', 3: 'yg', 4: 'ys', 5: 'kfy',
        6: 'pw', 7: 'yanyao_1', 8: 'yanyao_2', 9: 'paper_cup', 10: 'musai',
        11: 'carrot', 12: 'potato', 13: 'potato_s', 14: 'potato_black',
        15: 'cizhuan', 16: 'eluanshi_guang', 17: 'stone', 18: 'zhuankuai_bai',
        19: 'zhuankuai_red', 20: 'empty'
    }

    # 模型参数
    model_h = 320
    model_w = 320
    nl = 3
    na = 3
    stride = [8., 16., 32.]
    anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]
    anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(nl, -1, 2)

    # 打开摄像头
    video = 1  # 摄像头设备编号，1表示默认摄像头
    cap = cv2.VideoCapture(video)  # 视频捕获对象
    flag_det = False  # 检测开关

    while True:
        success, img0 = cap.read()  # 读取每一帧
        if success:
            if flag_det:  # 如果检测开启
                t1 = time.time()  # 记录推理前的时间
                det_boxes, scores, ids = infer_img(img0, net, model_h, model_w, nl, na, stride, anchor_grid,
                                                   thred_nms=0.4, thred_cond=0.5)  # 推理
                t2 = time.time()  # 记录推理后的时间

                # 绘制检测框和标签
                for box, score, id in zip(det_boxes, scores, ids):
                    label = '%s:%.2f' % (dic_labels[id], score)
                    plot_one_box(box.astype(np.int16), img0, color=(255, 0, 0), label=label, line_thickness=None)

                # 计算并显示FPS
                delta_time = t2 - t1
                if delta_time > 0:
                    str_FPS = "FPS: %.2f" % (1. / delta_time)
                else:
                    str_FPS = "FPS: inf"

                cv2.putText(img0, str_FPS, (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 3)

            cv2.imshow("video", img0)  # 显示图像
        key = cv2.waitKey(1) & 0xFF  # 等待键盘输入
        if key == ord('q'):  # 按q键退出
            break
        elif key & 0xFF == ord('s'):  # 按s键切换检测开关
            flag_det = not flag_det
            print(flag_det)

    cap.release()  # 释放视频捕获对象

代码中推理相关的函数做逐行详细中文注释和解释，帮助你从初学者的角度完全理解代码。我们重点讲解这几个核心函数：

---

1. infer_img() 推理主函数

def infer_img(img0, net, model_h, model_w, nl, na, stride, anchor_grid, thred_nms=0.4, thred_cond=0.5):

这是图像推理的主函数，完成从原始图像到预测结果的所有操作。

---

第一步：图像预处理

img = cv2.resize(img0, [model_w, model_h], interpolation=cv2.INTER_AREA)

• 将原始图像 img0 缩放成模型输入要求的大小（例如 320×320）。

• cv2.INTER_AREA 是一种图像插值方式，适合缩小图像时使用。

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

• OpenCV 读取图像是 BGR 顺序，而深度学习模型通常使用 RGB，因此这里需要转换颜色通道。

img = img.astype(np.float32) / 255.0

• 把图像的数据类型转为 float32，并将像素值从 [0, 255] 范围归一化到 [0, 1]，符合模型输入要求。

blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)

• OpenCV图像的格式是 (H, W, C)，而 PyTorch 模型（如YOLO）的输入是 (B, C, H, W)

• np.transpose(img, (2, 0, 1)) 把通道 C 移到第一个维度

• np.expand_dims(..., axis=0) 增加 batch 维度：变成 (1, 3, 320, 320)

---

第二步：模型推理

outs = net.run(None, {net.get_inputs()[0].name: blob})[0].squeeze(axis=0)

• 用 ONNX Runtime 推理：输入是 blob

• net.get_inputs()[0].name 得到模型输入的名字

• squeeze(axis=0) 把 batch 维度去掉，形状变成 (N, 85)，N 是预测框数量，85 是每个框的信息（x, y, w, h, conf, + 80类）

---

第三步：输出坐标解码

outs = cal_outputs(outs, nl, na, model_w, model_h, anchor_grid, stride)

• YOLO 的输出是相对 anchor + grid 编码的，需要转换为图像上的真实位置

• cal_outputs() 就是做这个解码变换的函数（后面详细讲）

---

第四步：后处理，获取检测框信息

img_h, img_w, _ = np.shape(img0)
boxes, confs, ids = post_process_opencv(outs, model_h, model_w, img_h, img_w, thred_nms, thred_cond)

• 将模型输出映射回原始图像尺寸

• 使用置信度阈值和 NMS 非极大值抑制删除重复框

• 得到最终的：

  • boxes: 框坐标

  • confs: 置信度

  • ids: 类别编号

---

2. cal_outputs() 坐标解码函数

def cal_outputs(outs, nl, na, model_w, model_h, anchor_grid, stride):

含义解释：

• outs: 模型输出，形状大致是 (N, 85)，前4列是框的位置

• nl: YOLO使用的输出层数量（3个：大中小目标）

• na: 每个特征层使用的 anchor 数（通常为 3）

• anchor_grid: 每层 anchor 的宽高尺寸

• stride: 每层特征图相对于原图的缩放倍数

grid = [np.zeros(1)] * nl

• 每一层都要生成网格坐标 grid，初始化为占位

for i in range(nl):
    h, w = int(model_w / stride[i]), int(model_h / stride[i])

• 计算第 i 层的特征图尺寸（如：320/8=40）

    length = int(na * h * w)

• 该层有多少个预测框

    if grid[i].shape[2:4] != (h, w):
        grid[i] = _make_grid(w, h)

• 如果还没有生成 grid，就调用 _make_grid() 创建形状为 (h*w, 2) 的网格点

    outs[row_ind:row_ind + length, 0:2] = ...
    outs[row_ind:row_ind + length, 2:4] = ...

• 对该层的所有框做位置矫正（中心点解码 + 宽高缩放）

• 用 grid 和 anchor 反算出真实坐标

---

3. post_process_opencv() 后处理函数

def post_process_opencv(outputs, model_h, model_w, img_h, img_w, thred_nms, thred_cond):

功能：

• 将模型输出映射回原始图像尺寸

• 提取类别信息

• 使用 OpenCV 的 cv2.dnn.NMSBoxes() 进行非极大值抑制，保留重要框

步骤：

conf = outputs[:, 4].tolist()         # 提取每个框的置信度
c_x = outputs[:, 0] / model_w * img_w
c_y = outputs[:, 1] / model_h * img_h
w = outputs[:, 2] / model_w * img_w
h = outputs[:, 3] / model_h * img_h

• 将中心点和尺寸从模型尺寸映射回原始图像尺寸

p_cls = outputs[:, 5:]
cls_id = np.argmax(p_cls, axis=1)

• 取得每个框的类别分数最大值（即分类结果）

p_x1 = c_x - w/2
p_y1 = c_y - h/2
p_x2 = c_x + w/2
p_y2 = c_y + h/2

• 把中心点转为左上角和右下角坐标 [x1, y1, x2, y2]

areas = np.concatenate((p_x1, p_y1, p_x2, p_y2), axis=-1)
ids = cv2.dnn.NMSBoxes(areas, conf, thred_cond, thred_nms)

• 用 NMS 去除重叠预测框

---