yolov5转onnx并推理
前言
该笔记记录了将训练后得出的yolov5s.pt模型转为.onnx模型并进行推理的过程。教程与大部分代码均来自这位作者的这篇博客。笔者仅对该代码进行修改和分析，在此十分感谢大佬们的帮助和分享！
代码
import onnxruntime
import numpy as np
import cv2
import os
import glob
CLASSES = ['smoke']
class YOLOV5():
def __init__(self,onnxpath):
self.onnx_session=onnxruntime.InferenceSession(onnxpath)
self.input_name=self.get_input_name()
self.output_name=self.get_output_name()
#-------------------------------------------------------
# 获取输入输出的名字
#-------------------------------------------------------
def get_input_name(self):
input_name=[]
for node in self.onnx_session.get_inputs():
input_name.append(node.name)
return input_name
def get_output_name(self):
output_name=[]
for node in self.onnx_session.get_outputs():
output_name.append(node.name)
return output_name
#-------------------------------------------------------
# 输入图像
#-------------------------------------------------------
def get_input_feed(self,img_tensor):
input_feed={}
for name in self.input_name:
input_feed[name]=img_tensor
return input_feed
#-------------------------------------------------------
# 1.cv2读取图像并resize
# 2.图像转BGR2RGB和HWC2CHW
# 3.图像归一化
# 4.图像增加维度
# 5.onnx_session 推理
#-------------------------------------------------------
def inference(self,img_path):
# cv2 读取图像
img=cv2.imread(img_path)
#进行resize
or_img=cv2.resize(img,(640,640))
#BGR2RGB和HWC2CHW（具体细节看代码解析）
img=or_img[:,:,::-1].transpose(2,0,1)
# 将图像数据类型转换为 float32
img=img.astype(dtype=np.float32)
# 将像素值归一化
img/=255.0
# 增加批次纬度
img=np.expand_dims(img,axis=0)
# 准备模型输入和推理
input_feed=self.get_input_feed(img)
pred=self.onnx_session.run(None,input_feed)[0]
# 返回模型的预测结果 pred 以及原始图像 or_img。
return pred,or_img
# 坐标转换
def xywh2xyxy(x):
# [x, y, w, h] to [x1, y1, x2, y2]
y = np.copy(x)
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
#dets: array [x,6] 6个值分别为x1,y1,x2,y2,score,class
#thresh: 阈值
def nms(dets, thresh):
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
#-------------------------------------------------------
# 计算框的面积
# 置信度从大到小排序
#-------------------------------------------------------
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
scores = dets[:, 4]
keep = []
index = scores.argsort()[::-1]
# 只要还有索引（即待处理的检测框）存在，循环就会继续执行。
while index.size > 0:
i = index[0]
keep.append(i)
#-------------------------------------------------------
# 计算相交面积
# 1.相交
# 2.不相交
#-------------------------------------------------------
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1)
h = np.maximum(0, y22 - y11 + 1)
overlaps = w * h
#-------------------------------------------------------
# 计算该框与其它框的IOU，去除掉重复的框，即IOU值大的框
# IOU小于thresh的框保留下来
#-------------------------------------------------------
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
idx = np.where(ious <= thresh)[0]
index = index[idx + 1]
return keep
def filter_box(org_box,conf_thres,iou_thres): #过滤掉无用的框
#-------------------------------------------------------
# 删除为1的维度
# 删除置信度小于conf_thres的BOX
#-------------------------------------------------------
org_box=np.squeeze(org_box)
conf = org_box[..., 4] > conf_thres
box = org_box[conf == True]
#-------------------------------------------------------
# 通过argmax获取置信度最大的类别
#-------------------------------------------------------
cls_cinf = box[..., 5:]
cls = []
for i in range(len(cls_cinf)):
cls.append(int(np.argmax(cls_cinf[i])))
all_cls = list(set(cls))
#-------------------------------------------------------
# 分别对每个类别进行过滤
# 1.将第6列元素替换为类别下标
# 2.xywh2xyxy 坐标转换
# 3.经过非极大抑制后输出的BOX下标
# 4.利用下标取出非极大抑制后的BOX
#-------------------------------------------------------
output = []
for i in range(len(all_cls)):
curr_cls = all_cls[i]
curr_cls_box = []
curr_out_box = []
for j in range(len(cls)):
if cls[j] == curr_cls:
box[j][5] = curr_cls
curr_cls_box.append(box[j][:6])
curr_cls_box = np.array(curr_cls_box)
# curr_cls_box_old = np.copy(curr_cls_box)
curr_cls_box = xywh2xyxy(curr_cls_box)
curr_out_box = nms(curr_cls_box,iou_thres)
for k in curr_out_box:
output.append(curr_cls_box[k])
output = np.array(output)
return output
def draw(image,box_data):
#-------------------------------------------------------
# 取整，方便画框
#-------------------------------------------------------
print(box_data)
boxes=box_data[...,:4].astype(np.int32)
scores=box_data[...,4]
classes=box_data[...,5].astype(np.int32)
for box, score, cl in zip(boxes, scores, classes):
top, left, right, bottom = box
print('class: {}, score: {}'.format(CLASSES[cl], score))
print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))
# cv2.rectangle 是 OpenCV 中用于在图像上绘制矩形的函数。这里用它来绘制边界框。
# (top, left) 和 (right, bottom) 分别是矩形（边界框）的左上角和右下角坐标。
# (255, 0, 0) 是颜色代码，这里表示蓝色。2 是线条的粗细。
cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
# cv2.putText 是 OpenCV 中用于在图像上绘制文本的函数。
# '{} {:.2f}'.format(CLASSES[cl], score) 生成标签文本，显示类别名称和置信度（保留两位小数）。
# (top, left) 是文本的起始坐标，通常设置在边界框的左上角。
# cv2.FONT_HERSHEY_SIMPLEX 指定字体类型。
# 0.6 是字体缩放比例。(0, 0, 255) 是字体颜色，这里表示红色。2 是线条的粗细。
cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
(top, left ),
cv2.FONT_HERSHEY_SIMPLEX,
0.6, (0, 0, 255), 2)
# cv2.imwrite('res.jpg',or_img) #保存该图片
cv2.imshow('Detected Image', or_img)
cv2.waitKey(0) # 等待用户按键
cv2.destroyAllWindows()
if __name__=="__main__":
onnx_path='testfile/best.onnx'
model=YOLOV5(onnx_path)
img_path = 'data/images'
img_paths = glob.glob(os.path.join(img_path, '*.jpg'))[:20]
for path in img_paths:
print(path)
output,or_img=model.inference(path)
outbox=filter_box(output,0.5,0.5)
if outbox != []:
draw(or_img,outbox)
这段代码是用于图像预处理、推理，以及获取预测结果的一系列操作。下面是对每一步的详细解释.
分析
YOLO5 class
图像格式转换:
img=or_img[:,:,::-1].transpose(2,0,1)
img=or_img[:,:,::-1]：从BGR到RGB的转换。在使用诸如OpenCV这样的库读取图像时，图像通常按BGR（蓝、绿、红）顺序存储颜色通道，这与大多数图像处理和深度学习库所期望的RGB（红、绿、蓝）顺序相反。为了适应这种差异，需要对图像颜色通道进行反转。这行代码利用NumPy的切片功能，[::-1] 表示以相反的顺序遍历数组的最后一个维度（即颜色通道）。这样，原始的BGR顺序就被颠倒为RGB顺序。or_img=or_img.transpose(2,0,1)高度-宽度-通道到通道-高度-宽度的转换。许多深度学习框架，如PyTorch和TensorFlow，预期输入数据的维度顺序为通道-高度-宽度（CHW），而不是传统的高度-宽度-通道（HWC）。为此，需要使用 .transpose() 方法来重排这些维度。这里，transpose(2, 0, 1) 指的是将原数组的第三个维度（颜色通道）移到前面，而将原来的第一个和第二个维度（高度和宽度）依次后移。这样，原始的HWC格式就被转换为CHW格式。
类型转换和归一化
img = img.astype(dtype=np.float32)
img /= 255.0
img.astype(dtype=np.float32)：将图像数据类型转换为float32，这是大多数深度学习框架处理图像数据的常用数据类型。img /= 255.0：将图像像素值从 [0, 255] 归一化到 [0, 1]。这种归一化有助于模型更好地处理输入数据，因为小的、归一化的输入值通常使模型训练更稳定。
增加批次维度:
img = np.expand_dims(img, axis=0)
np.expand_dims(img, axis=0)：在图像张量前添加一个额外的维度，将图像从单图像数组转换为批次大小为1的数组。这一步是必要的，因为大多数模型都是预期批次维度的输入（即使是单个图像也需要被视为批次大小为1的批次）。- 应用实例：在之前的步骤中，如果 img 经过
.transpose(2, 0, 1)调整后，假设它的形状是 (3, 高度, 宽度)。这表示一个单独的图像，其中3代表颜色通道。使用np.expand_dims(img, axis=0)将在索引0的位置（最前面）增加一个新的维度，形状变为 (1, 3, 高度, 宽度)。
准备模型输入
为了向模型提供正确的输入，通常需要构建一个名为 input_feed 的输入字典，这个字典将包含必要的输入数据。这一步是通过调用一个自定义方法 self.get_input_feed(img) 完成的，该方法的具体实现依赖于模型的具体输入需求：
input_feed = self.get_input_feed(img)
在这个方法中，img 通常是经过预处理的图像数据，input_feed 字典会根据模型的输入层期望的数据格式和类型，将图像数据封装成键值对。这里的键通常是输入层的名称，而值是对应的数据张量。
执行模型推理
一旦准备好输入字典，就可以使用 ONNX Runtime 的会话 (onnx_session) 来执行模型推理了。这通过 run 方法实现，该方法接收输出节点的名称（如果想获取所有输出可以传递 None）和输入字典：
pred = self.onnx_session.run(None, input_feed)[0]
self.onnx_session: 这是 onnxruntime.InferenceSession 的一个实例，负责加载 ONNX 模型并配置好所有必要的运行时参数。run方法: 此方法执行模型推理。它返回一个列表，其中包含对应于请求的输出节点的数据。传递 None 作为第一个参数意味着请求所有输出节点的数据。- 获取结果: 使用 [0] 从返回的列表中取出第一个元素，通常是主要的预测结果，这样 pred 就包含了模型的预测输出。
- 应用场景示例：假设有一个图像分类模型，其输入是处理过的图像数据，输出是分类的概率分布。在这种情况下：
- input_feed 将包括处理后的图像数据。
- self.onnx_session.run(None, input_feed)[0] 将执行推理并返回图像的分类概率。
- [0] 确保我们只获取主要输出，即分类概率。
坐标转换函数xywh2xyxy
函数 xywh2xyxy(x) 用于将图像的边界框坐标从一种格式转换到另一种格式。它将边界框的表示从中心坐标加宽度和高度（[x, y, w, h]）转换为边界框的左上角和右下角坐标（[x1, y1, x2, y2]）。
函数详解
输入 x 是一个 NumPy 数组，其中包含多个边界框的坐标，每个坐标由四个值组成：中心点的 x 坐标、中心点的 y 坐标、边界框的宽度、边界框的高度。这个函数按以下步骤处理这些坐标：
-
初始化y:
y = np.copy(x)：首先创建输入数组x的副本，以保持原始数据不变。所有的计算将在y上进行，这是输出数组。
-
计算左上角坐标 (x1, y1):
y[:, 0] = x[:, 0] - x[:, 2] / 2：这行代码计算边界框左上角的 x 坐标。它通过从边界框中心的 x 坐标 (x[:, 0]) 减去宽度的一半 (x[:, 2] / 2) 来完成。y[:, 1] = x[:, 1] - x[:, 3] / 2：这行代码计算边界框左上角的 y 坐标。它通过从边界框中心的 y 坐标 (x[:, 1]) 减去高度的一半 (x[:, 3] / 2) 来完成。
-
计算右下角坐标 (x2, y2):
y[:, 2] = x[:, 0] + x[:, 2] / 2：这行代码计算边界框右下角的 x 坐标。它通过在边界框中心的 x 坐标 (x[:, 0]) 上加上宽度的一半 (x[:, 2] / 2) 来完成。y[:, 3] = x[:, 1] + x[:, 3] / 2：这行代码计算边界框右下角的 y 坐标。它通过在边界框中心的 y 坐标 (x[:, 1]) 上加上高度的一半 (x[:, 3] / 2) 来完成。
-
返回结果:
return y：最后，函数返回y数组，其中包含了所有转换后的边界框坐标。
非极大值抑制函数nms
该函数 nms (Non-Maximum Suppression) 是计算机视觉和图像处理中常用的一种技术，用于去除冗余的边界框（bounding boxes），保留最佳的边界框，从而提高目标检测任务的准确性。函数的主要步骤如下：
输入参数
dets：一个二维数组，每一行表示一个检测到的对象的边界框，格式通常为 [x1, y1, x2, y2, score, class]，其中 (x1, y1) 是边界框左上角的坐标，(x2, y2) 是右下角的坐标，score 表示该边界框的置信度，class表示该边界框对应的目标种类。thresh：IOU（交并比）阈值，用于决定何时应该去除某个边界框。
主要步骤和逻辑
- 计算每个边界框的面积：
利用 (y2 - y1 + 1) 和 (x2 - x1 + 1) 计算每个边界框的面积，+1 是为了修正零长度的情况。 - 对置信度进行排序：
使用 argsort 对所有边界框的置信度 scores 进行降序排序，得到索引数组 index。 - 循环移除重叠的边界框：
- 使用循环结构遍历 index，从中选取置信度最高的边界框index[0]。
- 计算当前选中的边界框（置信度最高）与其他所有未处理的边界框的交集面积。
- 交集区域的确定应该是这样的：
- 左上角坐标（应是各个坐标的最大值）：
- x11: 交集左上角的 x 坐标应取当前框 x1[i] 和其他框 x1[index[1:]] 中的最大值。这是因为交集的左边界应在两个框的左边界中更靠右的位置，即较大的 x1 值。
- y11: 同理，交集左上角的 y 坐标应取 y1[i] 和 y1[index[1:]] 中的最大值，代表交集上边界在两个框的上边界中较低的位置。
- 右下角坐标（应是各个坐标的最小值）：
- x22: 交集右下角的 x 坐标应取当前框 x2[i] 和其他框 x2[index[1:]] 中的最小值。这是因为交集的右边界应在两个框的右边界中更靠左的位置，即较小的 x2 值。
- y22: 同理，交集右下角的 y 坐标应取 y2[i] 和 y2[index[1:]] 中的最小值，代表交集下边界在两个框的下边界中较高的位置。
- 左上角坐标（应是各个坐标的最大值）：
- 交集区域的确定应该是这样的：
- 通过计算交集的宽 w 和高 h，然后求乘积得到交集面积 overlaps。
- 根据交集面积和原有面积，计算IOU值 ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)。
- 根据IOU值与给定的阈值 thresh 判断哪些边界框应该保留：只保留那些IOU值小于 thresh 的边界框，即与当前选中的边界框重叠较少的边界框。
- 更新索引：
更新 index 以排除已经处理或决定去除的边界框，继续处理直到 index 为空。
这里注意，idx 是从 ious 数组中筛选得到的索引数组，代表了那些 IoU 小于阈值的框。但idx 数组中的索引是相对于 index[1:]（即除了当前最高置信度框之外的其他框）的。index[idx + 1] 这一操作实际上是将 idx 中的每个值增加 1（因为原始的 index 数组在计算 IoU 时被切片操作 index[1:] 排除了第一个元素，即当前最高置信度框）。这样可以将索引值转换回原始 index 数组的对应位置。 - 返回值
keep：一个列表，包含被保留的边界框的索引。
筛选预选框函数filter_box
函数 filter_box 是设计用来筛选和处理目标检测模型输出的边界框（bounding boxes），以提取置信度高且类别清晰的框，并对重叠较大的框应用非极大抑制（NMS）。下面详细解析这个函数的每个部分及其功能：
功能概述
- 输入:
org_box(原始的边界框数据),conf_thres(置信度阈值),iou_thres(用于NMS的IoU阈值)。 - 处理: 删除不符合置信度要求的框，对每个类别的框应用非极大抑制，以减少重叠。
- 输出: 经过筛选和NMS处理后的边界框。
详细步骤
-
删除为1的维度:
np.squeeze(org_box): 用于删除所有单维度条目，例如将形状为(1, N, 6)的数组压缩成(N, 6)。
-
删除置信度低于阈值的框:
conf = org_box[..., 4] > conf_thres: 生成一个布尔数组，标识哪些框的置信度高于conf_thres。box = org_box[conf == True]: 使用上面的布尔数组筛选出置信度高于阈值的框。
-
通过argmax获取置信度最大的类别:
cls_cinf = box[..., 5:]: 提取每个框的类别得分部分。cls = []: 初始化类别列表。for i in range(len(cls_cinf)): 遍历每个框，用argmax找到得分最高的类别索引，并添加到cls列表中。
-
类别去重:
all_cls = list(set(cls)): 从cls中提取唯一的类别列表。
-
分别对每个类别进行过滤:
- 对每个唯一类别进行循环，提取属于该类别的框。
- 将第6列元素替换为类别下标，准备进行NMS。
-
坐标转换和非极大抑制:
xywh2xyxy(curr_cls_box): 转换坐标格式，从中心宽高格式转换为左上角和右下角坐标格式。curr_out_box = nms(curr_cls_box, iou_thres): 对当前类别的框应用NMS，移除重叠过多的框。
-
收集最终输出:
- 通过NMS筛选后的框被添加到输出列表
output中。 output = np.array(output): 将输出列表转换为NumPy数组以便进一步处理或输出。
- 通过NMS筛选后的框被添加到输出列表
画图函数
函数 draw 主要负责在给定的图像上绘制边界框和对应的标签，用于可视化目标检测的结果。该函数接收两个参数：image 和 box_data，其中 image 是要绘制的图像，box_data 包含了边界框的位置、置信度和类别索引。下面逐步解析这个函数的各个部分及其功能：
参数解析
- image: 一个OpenCV图像对象，通常是通过
cv2.imread或相似方法读入的图像数据。 - box_data: 包含检测结果的数组，每一行代表一个检测到的对象，包括其边界框的坐标、置信度和类别索引。
功能详述
-
取整和类型转换:
boxes=box_data[...,:4].astype(np.int32): 从box_data中提取所有边界框的坐标并转换为整数。这里的...,:4意味着取数组的前四列，即边界框的x1, y1, x2, y2坐标。scores=box_data[...,4]: 提取每个边界框的置信度，位于第五列。classes=box_data[...,5].astype(np.int32): 提取每个边界框的类别索引，并转换为整数类型，位于第六列。
-
绘制边界框和标签:
- 使用
cv2.rectangle在图像上根据提供的边界框坐标绘制矩形。颜色(255, 0, 0)(蓝色) 和线条宽度 2 用于突出显示边界框。 - 使用
cv2.putText在图像上添加文本，显示每个边界框的类别和置信度。文本的起始位置设置在边界框的左上角，使用红色(0, 0, 255)和字体缩放比例 0.6 来增强可读性。
- 使用
-
交互式显示结果:
cv2.imshow('Detected Image', or_img): 显示处理后的图像，其中包括绘制的边界框和标签。cv2.waitKey(0): 等待用户按键响应，按任意键继续。cv2.destroyAllWindows(): 关闭所有OpenCV创建的窗口。
最后，看到这里如果觉得该笔记对您有用的话，可以点个小小的赞，或者点赞收藏关注一键三连ヾ(◍’౪`◍) ~ 谢谢！！
2万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
