1.上几篇文件,写了怎样标注数据,获取yolo数据,怎样训练数据。这篇文章写下,怎样使用yolo训练好的网络进行推理,怎样使用onnxruntime调用网络推理。
2.yolo模型训练后推理,主要分两种,一种是使用yolo官方提供的api推理,一种使用onnx模式推理。
3.使用yolo官方的api推理:
###把qimage图片转成mat函数
def CoverQImageToMat(self):
ptr_=self.CurrentImage.constBits() ###用pyqt加载的图片
###QImage内部是按照每个像素4个字节的方式组织数据的,即使最后一个alpha通道没有使用,也用0xff来填充
image_mat_=np.array(ptr_).reshape(self.CurrentImage.height(),self.CurrentImage.width(),4)
image_mat_rgb_=cv2.cvtColor(image_mat_,cv2.COLOR_RGBA2RGB)
return image_mat_rgb_
self.Model=YOLO(‘best.pt’)###加载模型
image_mat_ = self.CoverQImageToMat()###图片转换,把图片转成opencv的mat
results_=self.Model(image_mat_)###推理获取结果
for result_ in results_: #####循环结果
boxes_ = result_.boxes ####获取boxes
for box_ in boxes_: ####循环boxes
x_,y_,w_,h_=box_.xywh[0] ####获取识别到对象的左上角坐标,长,宽
c=box_.cls ####获取属于第几类
cls_=self.Model.names[int©] ####查询类名
self.BoxX.append(float(x_))
self.BoxY.append(float(y_))
self.BoxWidth.append(float(w_))
self.BoxHeight.append(float(h_))
self.BoxClass.append(cls_)
4.yolo官方提供的api推理,可以按照官方文档就可以推理,这里还在讲下导出的onnx模型后,怎么加载推理。
4.1self.Model.export(format=‘onnx’)###把yolo的模型导出成onnx格式
4.2 使用onnxruntime调用onnx推理。首先安装onnxruntime.lib,这里使用的版本是
onnxruntime-1.16.3-cp311-cp311-win_amd64.whl
onnx-1.15.0-cp311-cp311-win_amd64.whl
这两个文件可以到https://pypi.org/收索下载,使用pip离线安装
4.3 代码导入onnxruntime
import onnxruntime as ort
import onnx
self.OnnxModel = ort.InferenceSession(‘best.onnx’)###加载onnx模型
###推理函数
def UseOnnxInference(self):
if (self.CurrentImage.width() > 0) & (self.CurrentImage.height() > 0):
####清空结果缓存
self.BoxY.clear()
self.BoxX.clear()
self.BoxWidth.clear()
self.BoxHeight.clear()
self.BoxClass.clear()
#获取onnx模型输入点
inputs =self.OnnxModel.get_inputs()
len(inputs)
input_onnx_=inputs[0]
print("Name:",input_onnx_.name)##输入节点名称
print("Type:", input_onnx_.type)##数据格式
print("Shape:",input_onnx_.shape)###数据维度
image_mat_ = self.CoverQImageToMat() ###图片转换
print("ImageMatShape:",image_mat_.shape)####查看图片的维度
img_width_ = image_mat_.shape[1] ###获取图片输入输出
img_height_ = image_mat_.shape[0]###获取图片输入输出
target_image_height_ = 640###输入到onnx推理的图片大小
target_image_width_ = 640###输入到onnx推理的图片大小
scale_precentage_=min(target_image_height_/image_mat_.shape[0],target_image_width_/image_mat_.shape[1])
image_mat_=cv2.resize(image_mat_,None,fx=scale_precentage_,fy=scale_precentage_)#缩放成onnx推理的图片大小
image_np_=np.array(image_mat_)###图片转成np数组
print("ImageNpShape:", image_np_.shape)####查看下维度
image_np_=image_np_.transpose(2,0,1)##转成通道在前面的维度
print("ImageNpShape:", image_np_.shape)
image_np_=image_np_.reshape(1,3,640,640)##添加一个新维度
print(image_np_[0, 0, 0, 0])
print("ImageNpShape:", image_np_.shape)
image_np_ = image_np_.astype(np.float32)####把数据转成float32格式
image_np_=image_np_/255.0##数据归一化
print(image_np_[0, 0, 0, 0])
###获取输出点
outputs=self.OnnxModel.get_outputs()
output_onnx_=outputs[0]
print("Name:",output_onnx_.name)###输出节点名称
print("Type:",output_onnx_.type)###输出节点格式
print("Shape:",output_onnx_.shape)###输出节点维度
###运行推理
outputs=self.OnnxModel.run(["output0"],{"images":image_np_})
len(outputs)
###outputs处理
output_onnx_=outputs[0][0]
print("OutputOnxxShape:",output_onnx_.shape)
###转换 行列
output_onnx_=output_onnx_.transpose()
print("OutputOnxxShape:", output_onnx_.shape)
row=output_onnx_[0]
print(row)###输出下第一个特征
###获取多小行特征
rows_=output_onnx_.shape[0]
###获取boxs
boxs=[]
for i in range(rows_):
###获取最大概率 对应的index
prob_ = output_onnx_[i][4:].max()
if prob_>self.ConfidenceThres:
class_id_ = output_onnx_[i][4:].argmax()
###前是个特征代表 x中心 y中心 w宽 h高
xc_,yc_,w_,h_ = output_onnx_[i][:4]
###转换成x1 y1 x2 y2
x1_=(xc_-w_/2)/target_image_width_*img_width_
y1_=(yc_-h_/2)/target_image_height_*img_height_
x2_=(xc_+w_/2)/target_image_width_*img_width_
y2_=(yc_ + h_ / 2) / target_image_height_ * img_height_
boxs.append([x1_,y1_,x2_,y2_,class_id_,prob_])
###去掉重叠的box
boxs.sort(key=lambda x:[5],reverse=True)
result_box_=[]
lenght_box_=len(boxs)
while lenght_box_>0:
box_temp_=boxs[0].copy()
boxs_temp_ = []
##不是重叠的添加进去boxs_temp_
for box_ in boxs:
###计算intersection
box1_x1_,box1_y1_,box1_x2_,box1_y2_=box_temp_[:4]
box2_x1_, box2_y1_, box2_x2_, box2_y2_ = box_[:4]
distance1_=math.sqrt((box1_x1_-box2_x1_)**2+(box1_y1_-box2_y1_)**2)
distance2_=math.sqrt((box1_x2_-box2_x2_)**2+(box1_y2_-box2_y2_)**2)
if (distance1_
4.4这里是一个使用onnx检测安全帽子的demo
https://gitee.com/wenyuanmo/py-qt-load-yolo-or-onnx-check-hat/tree/master