当前位置： 首页 > news >正文

网站建设公司资讯,logo设计免费网址,天眼查企业查询公司,免费搭建wordpress本文分享YOLO11中#xff0c;从xxx.pt权重文件转为.onnx文件#xff0c;然后使用.onnx文件#xff0c;进行目标检测任务的模型推理。 用ONNX模型推理#xff0c;便于算法到开发板或芯片的部署。 备注#xff1a;本文是使用Python#xff0c;编写ONNX模型推理代码的 目…本文分享YOLO11中从xxx.pt权重文件转为.onnx文件然后使用.onnx文件进行目标检测任务的模型推理。 用ONNX模型推理便于算法到开发板或芯片的部署。 备注本文是使用Python编写ONNX模型推理代码的 目录 1、导出ONNX模型 2、目标检测——ONNX模型推理 2.1、ONNX模型推理-整体流程 2.2、预处理函数 2.3、后处理函数 2.4、检测效果可视化函数 2.5、YOLO11目标检测——ONNX模型推理完整代码  1、导出ONNX模型 首先我们训练好的模型生成xxx.pt权重文件 然后用下面代码导出ONNX模型简洁版 from ultralytics import YOLO# 加载一个模型路径为 YOLO 模型的 .pt 文件 model YOLO(runs/detect/train/weights/best.pt)# 导出模型格式为 ONNX model.export(formatonnx)运行代码后会在上面路径中生成best.onnx文件的 比如填写的路径是runs/detect/train/weights/best.pt那么在runs/detect/train/weights/目录中会生成与best.pt同名的onnx文件即best.onnx 上面代码示例是简单版如果需要更专业设置ONNX用下面版本的 YOLO11导出ONNX模型专业版 from ultralytics import YOLO# 加载一个模型路径为 YOLO 模型的 .pt 文件 model YOLO(runs/detect/train/weights/best.pt)# 导出模型设置多种参数 model.export(formatonnx, # 导出格式为 ONNXimgsz(640, 640), # 设置输入图像的尺寸kerasFalse, # 不导出为 Keras 格式optimizeFalse, # 不进行优化halfFalse, # 不启用 FP16 量化int8False, # 不启用 INT8 量化dynamicFalse, # 不启用动态输入尺寸simplifyTrue, # 简化 ONNX 模型opsetNone, # 使用最新的 opset 版本workspace4.0, # 为 TensorRT 优化设置最大工作区大小GiBnmsFalse, # 不添加 NMS非极大值抑制batch1 # 指定批处理大小 ) 对于model.export( )函数中各种参数说明 formatonnx指定导出模型的格式为 onnx。imgsz(640, 640)输入图像的尺寸设为 640x640。如果需要其他尺寸可以修改这个值。kerasFalse不导出为 Keras 格式的模型。optimizeFalse不应用 TorchScript 移动设备优化。halfFalse不启用 FP16半精度量化。int8False不启用 INT8 量化。dynamicFalse不启用动态输入尺寸。simplifyTrue简化模型以提升 ONNX 模型的性能。opsetNone使用默认的 ONNX opset 版本如果需要可以手动指定。workspace4.0为 TensorRT 优化指定最大工作空间大小为 4 GiB。nmsFalse不为 CoreML 导出添加非极大值抑制NMS。batch1设置批处理大小为 1。 参考官网文档https://docs.ultralytics.com/modes/export/#arguments 当然了YOLO11中不仅支持ONNX模型还支持下面表格中格式 支持的导出格式format参数值生成的模型示例model.export( )函数的参数PyTorch-yolo11n.pt-TorchScripttorchscriptyolo11n.torchscriptimgsz, optimize, batchONNXonnxyolo11n.onnximgsz, half, dynamic, simplify, opset, batchOpenVINOopenvinoyolo11n_openvino_model/imgsz, half, int8, batchTensorRTengineyolo11n.engineimgsz, half, dynamic, simplify, workspace, int8, batchCoreMLcoremlyolo11n.mlpackageimgsz, half, int8, nms, batchTF SavedModelsaved_modelyolo11n_saved_model/imgsz, keras, int8, batchTF GraphDefpbyolo11n.pbimgsz, batchTF Litetfliteyolo11n.tfliteimgsz, half, int8, batchTF Edge TPUedgetpuyolo11n_edgetpu.tfliteimgszTF.jstfjsyolo11n_web_model/imgsz, half, int8, batchPaddlePaddlepaddleyolo11n_paddle_model/imgsz, batchNCNNncnnyolo11n_ncnn_model/imgsz, half, batch 2、目标检测——ONNX模型推理 我们需要编写代码实现了一个使用 ONNXRuntime 执行 YOLOv11 检测模型推理的完整流程包含图像预处理、推理、后处理和可视化 。 需要编写的代码功能包括 预处理 读取输入图像通过 letterbox 填充图像自动调整图像尺寸使其符合模型的输入尺寸要求并归一化并转换为模型输入的格式 。 模型推理 使用 ONNXRuntime 加载YOLO11 模型并根据系统环境自动选择 CPU 或 GPU 执行推理。预处理后的图像输入到模型中模型返回预测的边界框、类别和分数等输出结果。 后处理 对模型输出的边界框、类别和分数进行处理首先根据置信度阈值过滤低置信度的检测结果。然后根据缩放比例将边界框映射回原图的尺寸应用非极大值抑制NMS去除重叠的边界框最终提取有效的检测结果。 可视化 在输入图像上绘制检测到的边界框、类别名称和置信度使用不同颜色区分不同类别的对象。可选择将带有检测结果的图像保存到文件中最终输出检测结果保存的路径。 2.1、ONNX模型推理-整体流程 首先编写一个用于运行YOLO11检测模型的推理类YOLO11 ONNX推理流程预处理 - 推理 - 后处理 class YOLO11:YOLO11 目标检测模型类用于处理推理和可视化。def init(self, onnx_model, input_image, confidence_thres, iou_thres):初始化 YOLO11 类的实例。参数onnx_model: ONNX 模型的路径。input_image: 输入图像的路径。confidence_thres: 用于过滤检测结果的置信度阈值。iou_thres: 非极大值抑制NMS的 IoU交并比阈值。self.onnx_model onnx_modelself.input_image input_imageself.confidence_thres confidence_thresself.iou_thres iou_thres# 加载类别名称self.classes CLASS_NAMES# 为每个类别生成一个颜色调色板self.color_palette np.random.uniform(0, 255, size(len(self.classes), 3))def main(self):# 使用 ONNX 模型创建推理会话自动选择CPU或GPUsession ort.InferenceSession(self.onnx_model, providers[CUDAExecutionProvider, CPUExecutionProvider] if ort.get_device() GPU else [CPUExecutionProvider],)# 打印模型的输入尺寸print(YOLO11 目标检测 ONNXRuntime)print(模型名称, self.onnx_model)# 获取模型的输入形状model_inputs session.get_inputs()input_shape model_inputs[0].shape self.input_width input_shape[2]self.input_height input_shape[3]print(f模型输入尺寸宽度 {self.input_width}, 高度 {self.input_height})# 预处理图像数据确保使用模型要求的尺寸 (640x640)img_data self.preprocess()# 使用预处理后的图像数据运行推理outputs session.run(None, {model_inputs[0].name: img_data})# 对输出进行后处理以获取输出图像return self.postprocess(self.img, outputs) # 输出图像 2.2、预处理函数 然后编写输入图像预处理函数preprocess 主要功能对输入的图像进行读取、颜色空间转换、尺寸调整、归一化等预处理操作并返回适合模型输入的图像数据。 使用 OpenCV 读取图像。将图像从 BGR 格式转换为 RGB 格式。使用 letterbox 函数保持图像的宽高比并填充图像使其符合模型要求的输入尺寸。归一化图像数据像素值除以 255.0。转换图像通道的维度符合 PyTorch 模型的输入格式通道优先。扩展图像维度以匹配模型输入的 batch 大小。 def preprocess(self):对输入图像进行预处理以便进行推理。返回image_data: 经过预处理的图像数据准备进行推理。# 使用 OpenCV 读取输入图像self.img cv2.imread(self.input_image)# 获取输入图像的高度和宽度self.img_height, self.img_width self.img.shape[:2]# 将图像颜色空间从 BGR 转换为 RGBimg cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)# 保持宽高比进行 letterbox 填充, 使用模型要求的输入尺寸img, self.ratio, (self.dw, self.dh) self.letterbox(img, new_shape(self.input_width, self.input_height))# 通过除以 255.0 来归一化图像数据image_data np.array(img) / 255.0# 将图像的通道维度移到第一维image_data np.transpose(image_data, (2, 0, 1)) # 通道优先# 扩展图像数据的维度以匹配模型输入的形状image_data np.expand_dims(image_data, axis0).astype(np.float32)# 返回预处理后的图像数据return image_data预处理函数preprocess会依赖letterbox函数。 letterbox函数功能将图像缩放并进行填充以保持宽高比最终将图像调整到指定的输入尺寸。 计算图像的缩放比例保证宽高比不变。根据缩放后的图像尺寸计算需要的填充量上下左右。对图像进行缩放和添加边框填充确保最终图像尺寸符合目标输入尺寸。返回调整后的图像、缩放比例及填充尺寸。 def letterbox(self, img, new_shape(640, 640), color(114, 114, 114), autoFalse, scaleFillFalse, scaleupTrue):将图像进行 letterbox 填充保持纵横比不变并缩放到指定尺寸。shape img.shape[:2] # 当前图像的宽高print(fOriginal image shape: {shape})if isinstance(new_shape, int):new_shape (new_shape, new_shape)# 计算缩放比例r min(new_shape[0] / shape[0], new_shape[1] / shape[1]) # 选择宽高中最小的缩放比if not scaleup: # 仅缩小不放大r min(r, 1.0)# 缩放后的未填充尺寸new_unpad (int(round(shape[1] * r)), int(round(shape[0] * r)))# 计算需要的填充dw, dh new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # 计算填充的尺寸dw / 2 # padding 均分dh / 2# 缩放图像if shape[::-1] ! new_unpad: # 如果当前图像尺寸不等于 new_unpad则缩放img cv2.resize(img, new_unpad, interpolationcv2.INTER_LINEAR)# 为图像添加边框以达到目标尺寸top, bottom int(round(dh)), int(round(dh))left, right int(round(dw)), int(round(dw))img cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, valuecolor)print(fFinal letterboxed image shape: {img.shape})return img, (r, r), (dw, dh) 2.3、后处理函数 再编写后处理函数postprocess 该代码的功能是对模型的输出结果进行后处理主要用于从目标检测模型的输出中提取边界框bounding boxes、置信度分数和类别ID并将这些检测结果绘制在输入图像上。  A、模型输出后处理 将模型的输出数据进行转置和压缩使其符合处理的预期格式。遍历模型的每一个检测结果提取检测框的坐标x, y, w, h、类别得分并通过阈值过滤掉低置信度的检测结果。计算检测框在原始图像上的实际位置考虑缩放和填充的影响。 B、缩放和填充调整  根据图像的缩放比例和填充量将检测框的坐标从网络输入尺寸如640x640调整回原始图像尺寸。 C、非极大值抑制NMS 使用非极大值抑制NMS算法对检测结果进行过滤删除重叠的检测框仅保留最高置信度的框。 D、绘制检测结果 使用 draw_detections 方法在输入图像上绘制边界框并标注检测类别和置信度。最终返回包含检测结果的图像图像上绘制了所有有效的检测框和对应的标签。 def postprocess(self, input_image, output):对模型输出进行后处理以提取边界框、分数和类别 ID。参数input_image (numpy.ndarray): 输入图像。output (numpy.ndarray): 模型的输出。返回numpy.ndarray: 包含检测结果的输入图像。# 转置并压缩输出以匹配预期形状outputs np.transpose(np.squeeze(output[0]))rows outputs.shape[0]boxes, scores, class_ids [], [], []# 计算缩放比例和填充ratio self.img_width / self.input_width, self.img_height / self.input_heightfor i in range(rows):classes_scores outputs[i][4:]max_score np.amax(classes_scores)if max_score self.confidence_thres:class_id np.argmax(classes_scores)x, y, w, h outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]# 将框调整到原始图像尺寸考虑缩放和填充x - self.dw # 移除填充y - self.dhx / self.ratio[0] # 缩放回原图y / self.ratio[1]w / self.ratio[0]h / self.ratio[1]left int(x - w / 2)top int(y - h / 2)width int(w)height int(h)boxes.append([left, top, width, height])scores.append(max_score)class_ids.append(class_id)indices cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)for i in indices:box boxes[i]score scores[i]class_id class_ids[i]self.draw_detections(input_image, box, score, class_id)return input_image 2.4、检测效果可视化函数 该代码的功能是根据检测结果在输入图像上绘制边界框和标签以可视化目标检测的结果。 def draw_detections(self, img, box, score, class_id):在输入图像上绘制检测到的边界框和标签。参数img: 用于绘制检测结果的输入图像。box: 检测到的边界框。score: 对应的检测分数。class_id: 检测到的目标类别 ID。返回None# 提取边界框的坐标x1, y1, w, h box# 获取类别对应的颜色color self.color_palette[class_id]# 在图像上绘制边界框cv2.rectangle(img, (int(x1), int(y1)), (int(x1 w), int(y1 h)), color, 2)# 创建包含类别名和分数的标签文本label f{self.classes[class_id]}: {score:.2f}# 计算标签文本的尺寸(label_width, label_height), _ cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)# 计算标签文本的位置label_x x1label_y y1 - 10 if y1 - 10 label_height else y1 10# 绘制填充的矩形作为标签文本的背景cv2.rectangle(img, (label_x, label_y - label_height), (label_x label_width, label_y label_height), color, cv2.FILLED)# 在图像上绘制标签文本cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA) 2.5、YOLO11目标检测——ONNX模型推理完整代码  完整代码如下所示

Ultralytics YOLO , AGPL-3.0 licenseimport argparse

import cv2 import numpy as np import onnxruntime as ort# 类外定义类别映射关系使用字典格式 CLASS_NAMES {0: class_name1, # 类别 0 名称1: class_name2, # 类别 1 名称2: class_name3 # 类别 1 名称# 可以添加更多类别… }class YOLO11:YOLO11 目标检测模型类用于处理推理和可视化。def init(self, onnx_model, input_image, confidence_thres, iou_thres):初始化 YOLO11 类的实例。参数onnx_model: ONNX 模型的路径。input_image: 输入图像的路径。confidence_thres: 用于过滤检测结果的置信度阈值。iou_thres: 非极大值抑制NMS的 IoU交并比阈值。self.onnx_model onnx_modelself.input_image input_imageself.confidence_thres confidence_thresself.iou_thres iou_thres# 加载类别名称self.classes CLASS_NAMES# 为每个类别生成一个颜色调色板self.color_palette np.random.uniform(0, 255, size(len(self.classes), 3))def preprocess(self):对输入图像进行预处理以便进行推理。返回image_data: 经过预处理的图像数据准备进行推理。# 使用 OpenCV 读取输入图像self.img cv2.imread(self.input_image)# 获取输入图像的高度和宽度self.img_height, self.img_width self.img.shape[:2]# 将图像颜色空间从 BGR 转换为 RGBimg cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)# 保持宽高比进行 letterbox 填充, 使用模型要求的输入尺寸img, self.ratio, (self.dw, self.dh) self.letterbox(img, new_shape(self.input_width, self.input_height))# 通过除以 255.0 来归一化图像数据image_data np.array(img) / 255.0# 将图像的通道维度移到第一维image_data np.transpose(image_data, (2, 0, 1)) # 通道优先# 扩展图像数据的维度以匹配模型输入的形状image_data np.expand_dims(image_data, axis0).astype(np.float32)# 返回预处理后的图像数据return image_datadef letterbox(self, img, new_shape(640, 640), color(114, 114, 114), autoFalse, scaleFillFalse, scaleupTrue):将图像进行 letterbox 填充保持纵横比不变并缩放到指定尺寸。shape img.shape[:2] # 当前图像的宽高print(fOriginal image shape: {shape})if isinstance(new_shape, int):new_shape (new_shape, new_shape)# 计算缩放比例r min(new_shape[0] / shape[0], new_shape[1] / shape[1]) # 选择宽高中最小的缩放比if not scaleup: # 仅缩小不放大r min(r, 1.0)# 缩放后的未填充尺寸new_unpad (int(round(shape[1] * r)), int(round(shape[0] * r)))# 计算需要的填充dw, dh new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # 计算填充的尺寸dw / 2 # padding 均分dh / 2# 缩放图像if shape[::-1] ! new_unpad: # 如果当前图像尺寸不等于 new_unpad则缩放img cv2.resize(img, new_unpad, interpolationcv2.INTER_LINEAR)# 为图像添加边框以达到目标尺寸top, bottom int(round(dh)), int(round(dh))left, right int(round(dw)), int(round(dw))img cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, valuecolor)print(fFinal letterboxed image shape: {img.shape})return img, (r, r), (dw, dh)def postprocess(self, input_image, output):对模型输出进行后处理以提取边界框、分数和类别 ID。参数input_image (numpy.ndarray): 输入图像。output (numpy.ndarray): 模型的输出。返回numpy.ndarray: 包含检测结果的输入图像。# 转置并压缩输出以匹配预期形状outputs np.transpose(np.squeeze(output[0]))rows outputs.shape[0]boxes, scores, class_ids [], [], []# 计算缩放比例和填充ratio self.img_width / self.input_width, self.img_height / self.input_heightfor i in range(rows):classes_scores outputs[i][4:]max_score np.amax(classes_scores)if max_score self.confidence_thres:class_id np.argmax(classes_scores)x, y, w, h outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]# 将框调整到原始图像尺寸考虑缩放和填充x - self.dw # 移除填充y - self.dhx / self.ratio[0] # 缩放回原图y / self.ratio[1]w / self.ratio[0]h / self.ratio[1]left int(x - w / 2)top int(y - h / 2)width int(w)height int(h)boxes.append([left, top, width, height])scores.append(max_score)class_ids.append(class_id)indices cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)for i in indices:box boxes[i]score scores[i]class_id class_ids[i]self.draw_detections(input_image, box, score, class_id)return input_imagedef draw_detections(self, img, box, score, class_id):在输入图像上绘制检测到的边界框和标签。参数img: 用于绘制检测结果的输入图像。box: 检测到的边界框。score: 对应的检测分数。class_id: 检测到的目标类别 ID。返回None# 提取边界框的坐标x1, y1, w, h box# 获取类别对应的颜色color self.color_palette[class_id]# 在图像上绘制边界框cv2.rectangle(img, (int(x1), int(y1)), (int(x1 w), int(y1 h)), color, 2)# 创建包含类别名和分数的标签文本label f{self.classes[class_id]}: {score:.2f}# 计算标签文本的尺寸(label_width, label_height), _ cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)# 计算标签文本的位置label_x x1label_y y1 - 10 if y1 - 10 label_height else y1 10# 绘制填充的矩形作为标签文本的背景cv2.rectangle(img, (label_x, label_y - label_height), (label_x label_width, label_y label_height), color, cv2.FILLED)# 在图像上绘制标签文本cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)def main(self):# 使用 ONNX 模型创建推理会话自动选择CPU或GPUsession ort.InferenceSession(self.onnx_model, providers[CUDAExecutionProvider, CPUExecutionProvider] if ort.get_device() GPU else [CPUExecutionProvider],)# 打印模型的输入尺寸print(YOLO11 目标检测 ONNXRuntime)print(模型名称, self.onnx_model)# 获取模型的输入形状model_inputs session.get_inputs()input_shape model_inputs[0].shape self.input_width input_shape[2]self.input_height input_shape[3]print(f模型输入尺寸宽度 {self.input_width}, 高度 {self.input_height})# 预处理图像数据确保使用模型要求的尺寸 (640x640)img_data self.preprocess()# 使用预处理后的图像数据运行推理outputs session.run(None, {model_inputs[0].name: img_data})# 对输出进行后处理以获取输出图像return self.postprocess(self.img, outputs) # 输出图像if name main:# 创建参数解析器以处理命令行参数parser argparse.ArgumentParser()parser.add_argument(–model, typestr, defaultruns/detect/train/weights/best.onnx, help输入你的 ONNX 模型路径。)parser.add_argument(–img, typestr, defaultrdatasets/test.jpg, help输入图像的路径。)parser.add_argument(–conf-thres, typefloat, default0.5, help置信度阈值)parser.add_argument(–iou-thres, typefloat, default0.45, helpNMS IoU 阈值)args parser.parse_args()# 使用指定的参数创建 YOLO11 类的实例detection YOLO11(args.model, args.img, args.conf_thres, args.iou_thres)# 执行目标检测并获取输出图像output_image detection.main()# 保存输出图像到文件cv2.imwrite(det_result_picture.jpg, output_image)print(图像已保存为 det_result_picture.jpg) 需要修改类别映射关系以及类别对应的颜色 比如定义两个类别car、person 示例代码

类外定义类别映射关系使用字典格式

CLASS_NAMES {0: car, # 类别 0 名称1: person # 类别 1 名称# 可以添加更多类别… } 运行代码打印信息 YOLO11 目标检测 ONNXRuntime 模型名称 runs/detect/train/weights/best.onnx 模型输入尺寸宽度 640, 高度 640 Original image shape: (398, 700) Final letterboxed image shape: (640, 640, 3) 图像已保存为 det_result_picture.jpg 可视化看一下分割效果保存名称是det_result_picture.jpg YOLO11相关文章推荐 一篇文章快速认识YOLO11 | 关键改进点 | 安装使用 | 模型训练和推理-CSDN博客 一篇文章快速认识 YOLO11 | 实例分割 | 模型训练 | 自定义数据集-CSDN博客 YOLO11模型推理 | 目标检测与跟踪 | 实例分割 | 关键点估计 | OBB旋转目标检测-CSDN博客 YOLO11模型训练 | 目标检测与跟踪 | 实例分割 | 关键点姿态估计-CSDN博客 YOLO11 实例分割 | 自动标注 | 预标注 | 标签格式转换 | 手动校正标签-CSDN博客 YOLO11 实例分割 | 导出ONNX模型 | ONNX模型推理-CSDN博客 分享完成欢迎大家多多点赞和收藏谢谢~