代码如下
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 = 0grid = [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 += lengthreturn 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, idsif __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 = 320model_w = 320nl = 3na = 3stride = [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)# 计算并显示FPSdelta_time = t2 - t1if 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键退出breakelif key & 0xFF == ord('s'): # 按s键切换检测开关flag_det = not flag_detprint(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 去除重叠预测框
)
