用ONNX模型推理#xff0c;便于算法到开发板或芯片的部署。

备注#xff1a;本文是使用Python#xff0c;编写ONNX模型推理代码的

目…本文分享YOLO11中从xxx.pt权重文件转为.onnx文件然后使用.onnx文件进行目标检测任务的模型推理。

1、导出ONNX模型

YOLO(runs/detect/train/weights/best.pt)#

导出模型格式为

model.export(formatonnx)运行代码后会在上面路径中生成best.onnx文件的

比如填写的路径是runs/detect/train/weights/best.pt那么在runs/detect/train/weights/目录中会生成与best.pt同名的onnx文件即best.onnx

YOLO11导出ONNX模型专业版

YOLO(runs/detect/train/weights/best.pt)#

导出模型设置多种参数

640x640。

如果需要其他尺寸可以修改这个值。

kerasFalse不导出为

Keras

量化。

dynamicFalse不启用动态输入尺寸。

simplifyTrue简化模型以提升

ONNX

参考官网文档https://docs.ultralytics.com/modes/export/#arguments

支持的导出格式format参数值生成的模型示例model.export(

)函数的参数PyTorch-yolo11n.pt-TorchScripttorchscriptyolo11n.torchscriptimgsz,

optimize,

batchONNXonnxyolo11n.onnximgsz,

half,

batchOpenVINOopenvinoyolo11n_openvino_model/imgsz,

half,

batchTensorRTengineyolo11n.engineimgsz,

half,

batchCoreMLcoremlyolo11n.mlpackageimgsz,

half,

SavedModelsaved_modelyolo11n_saved_model/imgsz,

keras,

TPUedgetpuyolo11n_edgetpu.tfliteimgszTF.jstfjsyolo11n_web_model/imgsz,

half,

batchPaddlePaddlepaddleyolo11n_paddle_model/imgsz,

batchNCNNncnnyolo11n_ncnn_model/imgsz,

half,

填充图像自动调整图像尺寸使其符合模型的输入尺寸要求并归一化并转换为模型输入的格式

模型推理

执行推理。

预处理后的图像输入到模型中模型返回预测的边界框、类别和分数等输出结果。

后处理

对模型输出的边界框、类别和分数进行处理首先根据置信度阈值过滤低置信度的检测结果。

然后根据缩放比例将边界框映射回原图的尺寸应用非极大值抑制NMS去除重叠的边界框最终提取有效的检测结果。

在输入图像上绘制检测到的边界框、类别名称和置信度使用不同颜色区分不同类别的对象。

可选择将带有检测结果的图像保存到文件中最终输出检测结果保存的路径。

2.1、ONNX模型推理-整体流程

input_imageself.confidence_thres

iou_thres#

为每个类别生成一个颜色调色板self.color_palette

np.random.uniform(0,

ort.InferenceSession(self.onnx_model,

providers[CUDAExecutionProvider,

ort.get_device()

session.get_inputs()input_shape

self.input_width

input_shape[2]self.input_height

{self.input_width},

主要功能对输入的图像进行读取、颜色空间转换、尺寸调整、归一化等预处理操作并返回适合模型输入的图像数据。

OpenCV

函数保持图像的宽高比并填充图像使其符合模型要求的输入尺寸。

归一化图像数据像素值除以

PyTorch

preprocess(self):对输入图像进行预处理以便进行推理。

返回image_data:

OpenCV

image_data预处理函数preprocess会依赖letterbox函数。

letterbox函数功能将图像缩放并进行填充以保持宽高比最终将图像调整到指定的输入尺寸。

计算图像的缩放比例保证宽高比不变。

根据缩放后的图像尺寸计算需要的填充量上下左右。

对图像进行缩放和添加边框填充确保最终图像尺寸符合目标输入尺寸。

返回调整后的图像、缩放比例及填充尺寸。

def

interpolationcv2.INTER_LINEAR)#

bottom

该代码的功能是对模型的输出结果进行后处理主要用于从目标检测模型的输出中提取边界框bounding

boxes、置信度分数和类别ID并将这些检测结果绘制在输入图像上。

A、模型输出后处理

将模型的输出数据进行转置和压缩使其符合处理的预期格式。

遍历模型的每一个检测结果提取检测框的坐标x,

h、类别得分并通过阈值过滤掉低置信度的检测结果。

计算检测框在原始图像上的实际位置考虑缩放和填充的影响。

B、缩放和填充调整

根据图像的缩放比例和填充量将检测框的坐标从网络输入尺寸如640x640调整回原始图像尺寸。

C、非极大值抑制NMS

使用非极大值抑制NMS算法对检测结果进行过滤删除重叠的检测框仅保留最高置信度的框。

D、绘制检测结果

方法在输入图像上绘制边界框并标注检测类别和置信度。

最终返回包含检测结果的图像图像上绘制了所有有效的检测框和对应的标签。

def

np.transpose(np.squeeze(output[0]))rows

scores,

height])scores.append(max_score)class_ids.append(class_id)indices

scores,

class_ids[i]self.draw_detections(input_image,

box,

该代码的功能是根据检测结果在输入图像上绘制边界框和标签以可视化目标检测的结果。

def

class_id):在输入图像上绘制检测到的边界框和标签。

参数img:

y1,

绘制填充的矩形作为标签文本的背景cv2.rectangle(img,

(label_x,

input_imageself.confidence_thres

iou_thres#

为每个类别生成一个颜色调色板self.color_palette

np.random.uniform(0,

preprocess(self):对输入图像进行预处理以便进行推理。

返回image_data:

OpenCV

interpolationcv2.INTER_LINEAR)#

bottom

np.transpose(np.squeeze(output[0]))rows

scores,

height])scores.append(max_score)class_ids.append(class_id)indices

scores,

class_ids[i]self.draw_detections(input_image,

box,

class_id):在输入图像上绘制检测到的边界框和标签。

参数img:

y1,

绘制填充的矩形作为标签文本的背景cv2.rectangle(img,

(label_x,

ort.InferenceSession(self.onnx_model,

providers[CUDAExecutionProvider,

ort.get_device()

session.get_inputs()input_shape

self.input_width

input_shape[2]self.input_height

{self.input_width},

argparse.ArgumentParser()parser.add_argument(--model,

typestr,

defaultruns/detect/train/weights/best.onnx,

ONNX

模型路径。

)parser.add_argument(--img,

typestr,

help输入图像的路径。

)parser.add_argument(--conf-thres,

typefloat,

help置信度阈值)parser.add_argument(--iou-thres,

typefloat,

保存输出图像到文件cv2.imwrite(det_result_picture.jpg,

比如定义两个类别car、person

runs/detect/train/weights/best.onnx

模型输入尺寸宽度

可视化看一下分割效果保存名称是det_result_picture.jpg

YOLO11相关文章推荐