YOLO11 实例分割 | 导出ONNX模型 | ONNX模型推理

本文分享YOLO11中，从xxx.pt权重文件转为.onnx文件，然后使用.onnx文件，进行实例分割任务的模型推理。

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

备注：本文是使用Python，编写ONNX模型推理代码的

目录

1、导出ONNX模型

2、实例分割——ONNX模型推理

2.1、整体ONNX推理流程代码

2.1、 图像预处理函数

2.3、 后处理函数

2.4、分割效果可视化函数

2.5、YOLO11实例分割——ONNX模型推理完整代码 

1、导出ONNX模型

首先我们训练好的模型，生成xxx.pt权重文件；

然后用下面代码，导出ONNX模型（简洁版）

  1. from ultralytics import YOLO
  2.  
  3. # 加载一个模型，路径为 YOLO 模型的 .pt 文件
  4. model = YOLO("runs/segment/train6/weights/best.pt")
  5.  
  6. # 导出模型，格式为 ONNX
  7. model.export(format="onnx")

运行代码后，会在上面路径中生成best.onnx文件的

  • 比如，填写的路径是："runs/segment/train6/weights/best.pt"
  • 那么在runs/segment/train6/weights/目录中，会生成与best.pt同名的onnx文件，即best.onnx

上面代码示例是简单版，如果需要更专业设置ONNX，用下面版本的

YOLO11导出ONNX模型（专业版）

  1. from ultralytics import YOLO
  2.  
  3. # 加载一个模型，路径为 YOLO 模型的 .pt 文件
  4. model = YOLO("runs/segment/train6/weights/best.pt")
  5.  
  6. # 导出模型，设置多种参数
  7. model.export(
  8.     format="onnx",      # 导出格式为 ONNX
  9.     imgsz=(640, 640),   # 设置输入图像的尺寸
  10.     keras=False,        # 不导出为 Keras 格式
  11.     optimize=False,     # 不进行优化
  12.     half=False,         # 不启用 FP16 量化
  13.     int8=False,         # 不启用 INT8 量化
  14.     dynamic=False,      # 不启用动态输入尺寸
  15.     simplify=True,      # 简化 ONNX 模型
  16.     opset=None,         # 使用最新的 opset 版本
  17.     workspace=4.0,      # 为 TensorRT 优化设置最大工作区大小（GiB）
  18.     nms=False,          # 不添加 NMS（非极大值抑制）
  19.     batch=1             # 指定批处理大小
  20. )

对于model.export( )函数中，各种参数说明：

  1. format="onnx"：指定导出模型的格式为 ONNX。
  2. imgsz=(640, 640)：输入图像的尺寸设为 640x640。如果需要其他尺寸可以修改这个值。
  3. keras=False：不导出为 Keras 格式的模型。
  4. optimize=False：不应用 TorchScript 移动设备优化。
  5. half=False：不启用 FP16（半精度）量化。
  6. int8=False：不启用 INT8 量化。
  7. dynamic=False：不启用动态输入尺寸。
  8. simplify=True：简化模型以提升 ONNX 模型的性能。
  9. opset=None：使用默认的 ONNX opset 版本，如果需要可以手动指定。
  10. workspace=4.0：为 TensorRT 优化指定最大工作空间大小为 4 GiB。
  11. nms=False：不为 CoreML 导出添加非极大值抑制（NMS）。
  12. batch=1：设置批处理大小为 1。

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

当然了，YOLO11中不但支持ONNX模型，还支持下面表格中格式

支持的导出格式format参数值生成的模型示例model.export( )函数的参数
PyTorch-yolo11n.pt-
TorchScripttorchscriptyolo11n.torchscriptimgsz, optimize, batch
ONNXonnxyolo11n.onnximgsz, half, dynamic, simplify, opset, batch
OpenVINOopenvinoyolo11n_openvino_model/imgsz, half, int8, batch
TensorRTengineyolo11n.engineimgsz, half, dynamic, simplify, workspace, int8, batch
CoreMLcoremlyolo11n.mlpackageimgsz, half, int8, nms, batch
TF SavedModelsaved_modelyolo11n_saved_model/imgsz, keras, int8, batch
TF GraphDefpbyolo11n.pbimgsz, batch
TF Litetfliteyolo11n.tfliteimgsz, half, int8, batch
TF Edge TPUedgetpuyolo11n_edgetpu.tfliteimgsz
TF.jstfjsyolo11n_web_model/imgsz, half, int8, batch
PaddlePaddlepaddleyolo11n_paddle_model/imgsz, batch
NCNNncnnyolo11n_ncnn_model/imgsz, half, batch

2、实例分割——ONNX模型推理

我们需要编写代码实现了一个使用 ONNXRuntime 执行 YOLOv11 分割模型推理的完整流程，包含图像预处理、推理、后处理和可视化 。

需要编写的代码功能包括：

  • 加载 YOLOv11 模型：代码使用 ONNXRuntime 加载和执行 YOLOv11 分割模型。
  • 支持不同尺寸图像的输入：代码对输入的图像进行预处理，保持原始宽高比，通过 letterbox 填充图像，使其符合模型的输入尺寸要求。
  • 模型推理：通过 ONNXRuntime 进行推理，生成预测的边界框、分割区域和掩膜。
  • 后处理：代码处理推理结果，应用置信度过滤和 NMS（非极大值抑制）来提取有效的检测结果，并对掩膜进行处理和转换为分割区域。
  • 可视化分割结果：代码支持可视化模型的输出结果，在图像上绘制检测到的分割区域、边界框和类别信息。支持将可视化结果保存到文件中。
  • 主函数：用户可以通过命令行参数指定 ONNX 模型路径、输入图像路径、置信度阈值和 IoU 阈值等参数，从而灵活配置模型推理。

2.1、整体ONNX推理流程代码

首先编写一个用于运行YOLO11分割模型的推理类：YOLO11Seg

ONNX推理流程：预处理 -> 推理 -> 后处理

  1. class YOLO11Seg:
  2.     def __init__(self, onnx_model):
  3.         # 创建 Ort 推理会话，选择 CPU 或 GPU 提供者
  4.         self.session = ort.InferenceSession(
  5.             onnx_model,
  6.             providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
  7.             if ort.get_device() == "GPU"
  8.             else ["CPUExecutionProvider"],
  9.         )
  10.         # 根据 ONNX 模型类型选择 Numpy 数据类型（支持 FP32 和 FP16）
  11.         self.ndtype = np.half if self.session.get_inputs()[0].type == "tensor(float16)" else np.single
  12.  
  13.         # 获取模型的输入宽度和高度（YOLO11-seg 只有一个输入）
  14.         self.model_height, self.model_width = [x.shape for x in self.session.get_inputs()][0][-2:]
  15.  
  16.         # 打印模型的输入尺寸
  17.         print("YOLO11 ? 实例分割 ONNXRuntime")
  18.         print("模型名称：", onnx_model)
  19.         print(f"模型输入尺寸：宽度 = {self.model_width}, 高度 = {self.model_height}")
  20.  
  21.         # 加载类别名称
  22.         self.classes = CLASS_NAMES
  23.  
  24.         # 加载类别对应的颜色
  25.         self.class_colors = CLASS_COLORS
  26.  
  27.  
  28.     def __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45, nm=32):
  29.         """
  30.         完整的推理流程：预处理 -> 推理 -> 后处理
  32.         Args:
  33.             im0 (Numpy.ndarray): 原始输入图像
  34.             conf_threshold (float): 置信度阈值
  35.             iou_threshold (float): NMS 中的 IoU 阈值
  36.             nm (int): 掩膜数量
  37.         Returns:
  38.             boxes (List): 边界框列表
  39.             segments (List): 分割区域列表
  40.             masks (np.ndarray): [N, H, W] 输出掩膜
  41.         """
  42.         # 图像预处理
  43.         im, ratio, (pad_w, pad_h) = self.preprocess(im0)
  44.  
  45.         # ONNX 推理
  46.         preds = self.session.run(None, {self.session.get_inputs()[0].name: im})
  47.  
  48.         # 后处理
  49.         boxes, segments, masks = self.postprocess(
  50.             preds,
  51.             im0=im0,
  52.             ratio=ratio,
  53.             pad_w=pad_w,
  54.             pad_h=pad_h,
  55.             conf_threshold=conf_threshold,
  56.             iou_threshold=iou_threshold,
  57.             nm=nm,
  58.         )
  59.         return boxes, segments, masks

2.1、 图像预处理函数

然后编写输入图像预处理函数：preprocess

  • 对不同尺寸的输入图像，使其适配YOLO模型的输入要求，同时保证图像的宽高比不变
  • 主要包含图像尺寸调整、填充、通道转换和归一化等操作
  1.     def preprocess(self, img):
  2.         """
  3.         图像预处理
  5.         Args:
  6.             img (Numpy.ndarray): 输入图像
  8.         Returns:
  9.             img_process (Numpy.ndarray): 处理后的图像
  10.             ratio (tuple): 宽高比例
  11.             pad_w (float): 宽度的填充
  12.             pad_h (float): 高度的填充
  13.         """
  14.         # 调整输入图像大小并使用 letterbox 填充
  15.         shape = img.shape[:2]  # 原始图像大小
  16.         new_shape = (self.model_height, self.model_width)
  17.         r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
  18.         ratio = r, r
  19.         new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  20.         pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # 填充宽高
  21.         if shape[::-1] != new_unpad:  # 调整图像大小
  22.             img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  23.         top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
  24.         left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
  25.         img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
  26.  
  27.         # 转换：HWC -> CHW -> BGR 转 RGB -> 除以 255 -> contiguous -> 添加维度
  28.         img = np.ascontiguousarray(np.einsum("HWC->CHW", img)[::-1], dtype=self.ndtype) / 255.0
  29.         img_process = img[None] if len(img.shape) == 3 else img
  30.         return img_process, ratio, (pad_w, pad_h)

2.3、 后处理函数

再编写YOLO11分割模型推理结果的后处理函数：postprocess

主要包括边界框和分割掩膜的处理

1、后处理框架 (postprocess)：

  • 获取预测输出的边界框和掩膜原型。
  • 过滤低置信度结果，结合边界框、置信度、类别和掩膜信息。
  • 应用非极大值抑制（NMS）去除重复框。
  • 将边界框缩放到原始图像的尺寸。
  • 处理并生成二值化的分割掩膜，并将其转换为分割区域。

2、掩膜处理：

  • 利用预测的掩膜原型和边界框生成最终的掩膜。
  • 缩放掩膜尺寸到原始图像大小，并进行裁剪以匹配边界框。

3、分割区域转换 (masks2segments)：

  • 将掩膜转换为分割区域的轮廓，用于进一步的分析或可视化。

4、掩膜缩放与裁剪：

  • 通过 scale_mask 和 crop_mask，将掩膜调整到与原始图像匹配的大小，并裁剪到边界框内部，生成最终的分割掩膜。
  1.     def postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold, nm=32):
  2.         """
  3.         推理后的结果后处理
  5.         Args:
  6.             preds (Numpy.ndarray): 来自 ONNX 的推理结果
  7.             im0 (Numpy.ndarray): [h, w, c] 原始输入图像
  8.             ratio (tuple): 宽高比例
  9.             pad_w (float): 宽度的填充
  10.             pad_h (float): 高度的填充
  11.             conf_threshold (float): 置信度阈值
  12.             iou_threshold (float): IoU 阈值
  13.             nm (int): 掩膜数量
  14.         Returns:
  15.             boxes (List): 边界框列表
  16.             segments (List): 分割区域列表
  17.             masks (np.ndarray): 掩膜数组
  18.         """
  19.         x, protos = preds[0], preds[1]  # 获取模型的两个输出：预测和原型
  20.  
  21.         # 转换维度
  22.         x = np.einsum("bcn->bnc", x)
  23.  
  24.         # 置信度过滤
  25.         x = x[np.amax(x[..., 4:-nm], axis=-1) > conf_threshold]
  26.  
  27.         # 合并边界框、置信度、类别和掩膜
  28.         x = np.c_[x[..., :4], np.amax(x[..., 4:-nm], axis=-1), np.argmax(x[..., 4:-nm], axis=-1), x[..., -nm:]]
  29.  
  30.         # NMS 过滤
  31.         x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]
  32.  
  33.         # 解析并返回结果
  34.         if len(x) > 0:
  35.             # 边界框格式转换：从 cxcywh -> xyxy
  36.             x[..., [0, 1]] -= x[..., [2, 3]] / 2
  37.             x[..., [2, 3]] += x[..., [0, 1]]
  38.  
  39.             # 缩放边界框，使其与原始图像尺寸匹配
  40.             x[..., :4] -= [pad_w, pad_h, pad_w, pad_h]
  41.             x[..., :4] /= min(ratio)
  42.  
  43.             # 限制边界框在图像边界内
  44.             x[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])
  45.             x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])
  46.  
  47.             # 处理掩膜
  48.             masks = self.process_mask(protos[0], x[:, 6:], x[:, :4], im0.shape)
  49.  
  50.             # 将掩膜转换为分割区域
  51.             segments = self.masks2segments(masks)
  52.             return x[..., :6], segments, masks  # 返回边界框、分割区域和掩膜
  53.         else:
  54.             return [], [], []
  55.  
  56.     @staticmethod
  57.     def masks2segments(masks):
  58.         """
  59.         将掩膜转换为分割区域
  61.         Args:
  62.             masks (numpy.ndarray): 模型输出的掩膜，形状为 (n, h, w)
  63.         Returns:
  64.             segments (List): 分割区域的列表
  65.         """
  66.         segments = []
  67.         for x in masks.astype("uint8"):
  68.             c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[0]  # 找到轮廓
  69.             if c:
  70.                 c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
  71.             else:
  72.                 c = np.zeros((0, 2))  # 如果没有找到分割区域，返回空数组
  73.             segments.append(c.astype("float32"))
  74.         return segments
  75.  
  76.     @staticmethod
  77.     def crop_mask(masks, boxes):
  78.         """
  79.         裁剪掩膜，使其与边界框对齐
  81.         Args:
  82.             masks (Numpy.ndarray): [n, h, w] 掩膜数组
  83.             boxes (Numpy.ndarray): [n, 4] 边界框
  84.         Returns:
  85.             (Numpy.ndarray): 裁剪后的掩膜
  86.         """
  87.         n, h, w = masks.shape
  88.         x1, y1, x2, y2 = np.split(boxes[:, :, None], 4, 1)
  89.         r = np.arange(w, dtype=x1.dtype)[None, None, :]
  90.         c = np.arange(h, dtype=x1.dtype)[None, :, None]
  91.         return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
  92.  
  93.     def process_mask(self, protos, masks_in, bboxes, im0_shape):
  94.         """
  95.         处理模型输出的掩膜
  97.         Args:
  98.             protos (numpy.ndarray): [mask_dim, mask_h, mask_w] 掩膜原型
  99.             masks_in (numpy.ndarray): [n, mask_dim] 掩膜数量
  100.             bboxes (numpy.ndarray): 缩放到原始图像尺寸的边界框
  101.             im0_shape (tuple): 原始输入图像的尺寸 (h,w,c)
  102.         Returns:
  103.             (numpy.ndarray): 处理后的掩膜
  104.         """
  105.         c, mh, mw = protos.shape
  106.         masks = np.matmul(masks_in, protos.reshape((c, -1))).reshape((-1, mh, mw)).transpose(1, 2, 0)  # HWN
  107.         masks = np.ascontiguousarray(masks)
  108.         masks = self.scale_mask(masks, im0_shape)  # 将掩膜从 P3 尺寸缩放到原始输入图像大小
  109.         masks = np.einsum("HWN -> NHW", masks)  # HWN -> NHW
  110.         masks = self.crop_mask(masks, bboxes)  # 裁剪掩膜
  111.         return np.greater(masks, 0.5)  # 返回二值化后的掩膜
  112.  
  113.     @staticmethod
  114.     def scale_mask(masks, im0_shape, ratio_pad=None):
  115.         """
  116.         将掩膜缩放至原始图像大小
  118.         Args:
  119.             masks (np.ndarray): 缩放和填充后的掩膜
  120.             im0_shape (tuple): 原始图像大小
  121.             ratio_pad (tuple): 填充与原始图像的比例
  122.         Returns:
  123.             masks (np.ndarray): 缩放后的掩膜
  124.         """
  125.         im1_shape = masks.shape[:2]
  126.         if ratio_pad is None:  # 计算比例
  127.             gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1])  # 比例
  128.             pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2  # 填充
  129.         else:
  130.             pad = ratio_pad[1]
  131.  
  132.         # 计算掩膜的边界
  133.         top, left = int(round(pad[1] - 0.1)), int(round(pad[0] - 0.1))  # y, x
  134.         bottom, right = int(round(im1_shape[0] - pad[1] + 0.1)), int(round(im1_shape[1] - pad[0] + 0.1))
  135.         if len(masks.shape) < 2:
  136.             raise ValueError(f'"len of masks shape" 应该是 2 或 3，但得到 {len(masks.shape)}')
  137.         masks = masks[top:bottom, left:right]
  138.         masks = cv2.resize(
  139.             masks, (im0_shape[1], im0_shape[0]), interpolation=cv2.INTER_LINEAR
  140.         )  # 使用 INTER_LINEAR 插值调整大小
  141.         if len(masks.shape) == 2:
  142.             masks = masks[:, :, None]
  143.         return masks

2.4、分割效果可视化函数

该代码的核心功能是绘制和可视化检测和分割结果

  1.     def draw_and_visualize(self, im, bboxes, segments, vis=False, save=True):
  2.         """
  3.         绘制和可视化结果
  5.         Args:
  6.             im (np.ndarray): 原始图像，形状为 [h, w, c]
  7.             bboxes (numpy.ndarray): [n, 4]，n 是边界框数量
  8.             segments (List): 分割区域的列表
  9.             vis (bool): 是否使用 OpenCV 显示图像
  10.             save (bool): 是否保存带注释的图像
  11.         Returns:
  12.             None
  13.         """
  14.         # 创建图像副本
  15.         im_canvas = im.copy()
  16.  
  17.         for (*box, conf, cls_), segment in zip(bboxes, segments):
  18.             # 获取类别对应的颜色
  19.             color = self.get_color_for_class(int(cls_))
  20.  
  21.             # 绘制轮廓和填充掩膜
  22.             # cv2.polylines(im, np.int32([segment]), True, (255, 255, 255), 2)  # 绘制白色边框
  23.             cv2.fillPoly(im_canvas, np.int32([segment]), color)  # 使用类别对应的颜色填充多边形
  24.  
  25.             # 绘制边界框
  26.             cv2.rectangle(im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 1, cv2.LINE_AA)
  27.             # 在图像上绘制类别名称和置信度
  28.             cv2.putText(im, f"{self.classes[cls_]}: {conf:.3f}", (int(box[0]), int(box[1] - 9)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)
  29.  
  30.         # 将图像和绘制的多边形混合
  31.         im = cv2.addWeighted(im_canvas, 0.3, im, 0.7, 0)
  32.  
  33.         # 显示图像
  34.         if vis:
  35.             cv2.imshow("seg_result_picture", im)
  36.             cv2.waitKey(0)
  37.             cv2.destroyAllWindows()
  38.  
  39.         # 保存图像
  40.         if save:
  41.             cv2.imwrite("seg_result_picture.jpg", im)

2.5、YOLO11实例分割——ONNX模型推理完整代码 

完整代码，如下所示：

  1. # Ultralytics YOLO ?, AGPL-3.0 license
  2. """ 
  3. YOLO11 分割模型 ONNXRuntime 
  4.     功能1: 支持不用尺寸图像的输入
  5.     功能2: 支持可视化分割结果
  6. """
  7.  
  8. import argparse
  9. import cv2
  10. import numpy as np
  11. import onnxruntime as ort
  12.  
  13. # 类外定义类别映射关系，使用字典格式
  14. CLASS_NAMES = {
  15.     0: 'class_name1',  # 类别 0 名称
  16.     1: 'class_name2'   # 类别 1 名称
  17.                        # 可以添加更多类别...
  18. }
  19.  
  20. # 定义类别对应的颜色，格式为 (R, G, B)
  21. CLASS_COLORS = {
  22.     0: (255, 255, 0),   # 类别 0 的颜色为青黄色
  23.     1: (255, 0, 0)      # 类别 1 的颜色为红色
  24.                         # 可以为其他类别指定颜色...
  25. }
  26.  
  27. class YOLO11Seg:
  28.     def __init__(self, onnx_model):
  29.         # 创建 Ort 推理会话，选择 CPU 或 GPU 提供者
  30.         self.session = ort.InferenceSession(
  31.             onnx_model,
  32.             providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
  33.             if ort.get_device() == "GPU"
  34.             else ["CPUExecutionProvider"],
  35.         )
  36.         # 根据 ONNX 模型类型选择 Numpy 数据类型（支持 FP32 和 FP16）
  37.         self.ndtype = np.half if self.session.get_inputs()[0].type == "tensor(float16)" else np.single
  38.  
  39.         # 获取模型的输入宽度和高度（YOLO11-seg 只有一个输入）
  40.         self.model_height, self.model_width = [x.shape for x in self.session.get_inputs()][0][-2:]
  41.  
  42.         # 打印模型的输入尺寸
  43.         print("YOLO11 ? 实例分割 ONNXRuntime")
  44.         print("模型名称：", onnx_model)
  45.         print(f"模型输入尺寸：宽度 = {self.model_width}, 高度 = {self.model_height}")
  46.  
  47.         # 加载类别名称
  48.         self.classes = CLASS_NAMES
  49.  
  50.         # 加载类别对应的颜色
  51.         self.class_colors = CLASS_COLORS
  52.  
  53.     def get_color_for_class(self, class_id):
  54.         return self.class_colors.get(class_id, (255, 255, 255))  # 如果没有找到类别颜色，返回白色
  55.  
  56.     def __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45, nm=32):
  57.         """
  58.         完整的推理流程：预处理 -> 推理 -> 后处理
  60.         Args:
  61.             im0 (Numpy.ndarray): 原始输入图像
  62.             conf_threshold (float): 置信度阈值
  63.             iou_threshold (float): NMS 中的 IoU 阈值
  64.             nm (int): 掩膜数量
  65.         Returns:
  66.             boxes (List): 边界框列表
  67.             segments (List): 分割区域列表
  68.             masks (np.ndarray): [N, H, W] 输出掩膜
  69.         """
  70.         # 图像预处理
  71.         im, ratio, (pad_w, pad_h) = self.preprocess(im0)
  72.  
  73.         # ONNX 推理
  74.         preds = self.session.run(None, {self.session.get_inputs()[0].name: im})
  75.  
  76.         # 后处理
  77.         boxes, segments, masks = self.postprocess(
  78.             preds,
  79.             im0=im0,
  80.             ratio=ratio,
  81.             pad_w=pad_w,
  82.             pad_h=pad_h,
  83.             conf_threshold=conf_threshold,
  84.             iou_threshold=iou_threshold,
  85.             nm=nm,
  86.         )
  87.         return boxes, segments, masks
  88.  
  89.     def preprocess(self, img):
  90.         """
  91.         图像预处理
  93.         Args:
  94.             img (Numpy.ndarray): 输入图像
  96.         Returns:
  97.             img_process (Numpy.ndarray): 处理后的图像
  98.             ratio (tuple): 宽高比例
  99.             pad_w (float): 宽度的填充
  100.             pad_h (float): 高度的填充
  101.         """
  102.         # 调整输入图像大小并使用 letterbox 填充
  103.         shape = img.shape[:2]  # 原始图像大小
  104.         new_shape = (self.model_height, self.model_width)
  105.         r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
  106.         ratio = r, r
  107.         new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  108.         pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # 填充宽高
  109.         if shape[::-1] != new_unpad:  # 调整图像大小
  110.             img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  111.         top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
  112.         left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
  113.         img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
  114.  
  115.         # 转换：HWC -> CHW -> BGR 转 RGB -> 除以 255 -> contiguous -> 添加维度
  116.         img = np.ascontiguousarray(np.einsum("HWC->CHW", img)[::-1], dtype=self.ndtype) / 255.0
  117.         img_process = img[None] if len(img.shape) == 3 else img
  118.         return img_process, ratio, (pad_w, pad_h)
  119.  
  120.     def postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold, nm=32):
  121.         """
  122.         推理后的结果后处理
  124.         Args:
  125.             preds (Numpy.ndarray): 来自 ONNX 的推理结果
  126.             im0 (Numpy.ndarray): [h, w, c] 原始输入图像
  127.             ratio (tuple): 宽高比例
  128.             pad_w (float): 宽度的填充
  129.             pad_h (float): 高度的填充
  130.             conf_threshold (float): 置信度阈值
  131.             iou_threshold (float): IoU 阈值
  132.             nm (int): 掩膜数量
  133.         Returns:
  134.             boxes (List): 边界框列表
  135.             segments (List): 分割区域列表
  136.             masks (np.ndarray): 掩膜数组
  137.         """
  138.         x, protos = preds[0], preds[1]  # 获取模型的两个输出：预测和原型
  139.  
  140.         # 转换维度
  141.         x = np.einsum("bcn->bnc", x)
  142.  
  143.         # 置信度过滤
  144.         x = x[np.amax(x[..., 4:-nm], axis=-1) > conf_threshold]
  145.  
  146.         # 合并边界框、置信度、类别和掩膜
  147.         x = np.c_[x[..., :4], np.amax(x[..., 4:-nm], axis=-1), np.argmax(x[..., 4:-nm], axis=-1), x[..., -nm:]]
  148.  
  149.         # NMS 过滤
  150.         x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]
  151.  
  152.         # 解析并返回结果
  153.         if len(x) > 0:
  154.             # 边界框格式转换：从 cxcywh -> xyxy
  155.             x[..., [0, 1]] -= x[..., [2, 3]] / 2
  156.             x[..., [2, 3]] += x[..., [0, 1]]
  157.  
  158.             # 缩放边界框，使其与原始图像尺寸匹配
  159.             x[..., :4] -= [pad_w, pad_h, pad_w, pad_h]
  160.             x[..., :4] /= min(ratio)
  161.  
  162.             # 限制边界框在图像边界内
  163.             x[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])
  164.             x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])
  165.  
  166.             # 处理掩膜
  167.             masks = self.process_mask(protos[0], x[:, 6:], x[:, :4], im0.shape)
  168.  
  169.             # 将掩膜转换为分割区域
  170.             segments = self.masks2segments(masks)
  171.             return x[..., :6], segments, masks  # 返回边界框、分割区域和掩膜
  172.         else:
  173.             return [], [], []
  174.  
  175.     @staticmethod
  176.     def masks2segments(masks):
  177.         """
  178.         将掩膜转换为分割区域
  180.         Args:
  181.             masks (numpy.ndarray): 模型输出的掩膜，形状为 (n, h, w)
  182.         Returns:
  183.             segments (List): 分割区域的列表
  184.         """
  185.         segments = []
  186.         for x in masks.astype("uint8"):
  187.             c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[0]  # 找到轮廓
  188.             if c:
  189.                 c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
  190.             else:
  191.                 c = np.zeros((0, 2))  # 如果没有找到分割区域，返回空数组
  192.             segments.append(c.astype("float32"))
  193.         return segments
  194.  
  195.     @staticmethod
  196.     def crop_mask(masks, boxes):
  197.         """
  198.         裁剪掩膜，使其与边界框对齐
  200.         Args:
  201.             masks (Numpy.ndarray): [n, h, w] 掩膜数组
  202.             boxes (Numpy.ndarray): [n, 4] 边界框
  203.         Returns:
  204.             (Numpy.ndarray): 裁剪后的掩膜
  205.         """
  206.         n, h, w = masks.shape
  207.         x1, y1, x2, y2 = np.split(boxes[:, :, None], 4, 1)
  208.         r = np.arange(w, dtype=x1.dtype)[None, None, :]
  209.         c = np.arange(h, dtype=x1.dtype)[None, :, None]
  210.         return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
  211.  
  212.     def process_mask(self, protos, masks_in, bboxes, im0_shape):
  213.         """
  214.         处理模型输出的掩膜
  216.         Args:
  217.             protos (numpy.ndarray): [mask_dim, mask_h, mask_w] 掩膜原型
  218.             masks_in (numpy.ndarray): [n, mask_dim] 掩膜数量
  219.             bboxes (numpy.ndarray): 缩放到原始图像尺寸的边界框
  220.             im0_shape (tuple): 原始输入图像的尺寸 (h,w,c)
  221.         Returns:
  222.             (numpy.ndarray): 处理后的掩膜
  223.         """
  224.         c, mh, mw = protos.shape
  225.         masks = np.matmul(masks_in, protos.reshape((c, -1))).reshape((-1, mh, mw)).transpose(1, 2, 0)  # HWN
  226.         masks = np.ascontiguousarray(masks)
  227.         masks = self.scale_mask(masks, im0_shape)  # 将掩膜从 P3 尺寸缩放到原始输入图像大小
  228.         masks = np.einsum("HWN -> NHW", masks)  # HWN -> NHW
  229.         masks = self.crop_mask(masks, bboxes)  # 裁剪掩膜
  230.         return np.greater(masks, 0.5)  # 返回二值化后的掩膜
  231.  
  232.     @staticmethod
  233.     def scale_mask(masks, im0_shape, ratio_pad=None):
  234.         """
  235.         将掩膜缩放至原始图像大小
  237.         Args:
  238.             masks (np.ndarray): 缩放和填充后的掩膜
  239.             im0_shape (tuple): 原始图像大小
  240.             ratio_pad (tuple): 填充与原始图像的比例
  241.         Returns:
  242.             masks (np.ndarray): 缩放后的掩膜
  243.         """
  244.         im1_shape = masks.shape[:2]
  245.         if ratio_pad is None:  # 计算比例
  246.             gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1])  # 比例
  247.             pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2  # 填充
  248.         else:
  249.             pad = ratio_pad[1]
  250.  
  251.         # 计算掩膜的边界
  252.         top, left = int(round(pad[1] - 0.1)), int(round(pad[0] - 0.1))  # y, x
  253.         bottom, right = int(round(im1_shape[0] - pad[1] + 0.1)), int(round(im1_shape[1] - pad[0] + 0.1))
  254.         if len(masks.shape) < 2:
  255.             raise ValueError(f'"len of masks shape" 应该是 2 或 3，但得到 {len(masks.shape)}')
  256.         masks = masks[top:bottom, left:right]
  257.         masks = cv2.resize(
  258.             masks, (im0_shape[1], im0_shape[0]), interpolation=cv2.INTER_LINEAR
  259.         )  # 使用 INTER_LINEAR 插值调整大小
  260.         if len(masks.shape) == 2:
  261.             masks = masks[:, :, None]
  262.         return masks
  263.  
  264.     def draw_and_visualize(self, im, bboxes, segments, vis=False, save=True):
  265.         """
  266.         绘制和可视化结果
  268.         Args:
  269.             im (np.ndarray): 原始图像，形状为 [h, w, c]
  270.             bboxes (numpy.ndarray): [n, 4]，n 是边界框数量
  271.             segments (List): 分割区域的列表
  272.             vis (bool): 是否使用 OpenCV 显示图像
  273.             save (bool): 是否保存带注释的图像
  274.         Returns:
  275.             None
  276.         """
  277.         # 创建图像副本
  278.         im_canvas = im.copy()
  279.  
  280.         for (*box, conf, cls_), segment in zip(bboxes, segments):
  281.             # 获取类别对应的颜色
  282.             color = self.get_color_for_class(int(cls_))
  283.  
  284.             # 绘制轮廓和填充掩膜
  285.             # cv2.polylines(im, np.int32([segment]), True, (255, 255, 255), 2)  # 绘制白色边框
  286.             cv2.fillPoly(im_canvas, np.int32([segment]), color)  # 使用类别对应的颜色填充多边形
  287.  
  288.             # 绘制边界框
  289.             cv2.rectangle(im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 1, cv2.LINE_AA)
  290.             # 在图像上绘制类别名称和置信度
  291.             cv2.putText(im, f"{self.classes[cls_]}: {conf:.3f}", (int(box[0]), int(box[1] - 9)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)
  292.  
  293.         # 将图像和绘制的多边形混合
  294.         im = cv2.addWeighted(im_canvas, 0.3, im, 0.7, 0)
  295.  
  296.         # 显示图像
  297.         if vis:
  298.             cv2.imshow("seg_result_picture", im)
  299.             cv2.waitKey(0)
  300.             cv2.destroyAllWindows()
  301.  
  302.         # 保存图像
  303.         if save:
  304.             cv2.imwrite("seg_result_picture.jpg", im)
  305.  
  306. if __name__ == "__main__":
  307.     # 创建命令行参数解析器
  308.     parser = argparse.ArgumentParser()
  309.     parser.add_argument("--model", type=str, default=r"weights/seg_best_exp6.onnx" , help="ONNX 模型路径")
  310.     parser.add_argument("--source", type=str, default=r"datasets/Mix_seg_point_offer_num53/images/point_offer_20240930_14_Image.png", help="输入图像路径")
  311.     parser.add_argument("--conf", type=float, default=0.6, help="置信度阈值")
  312.     parser.add_argument("--iou", type=float, default=0.45, help="NMS 的 IoU 阈值")
  313.     args = parser.parse_args()
  314.  
  315.     # 加载模型
  316.     model = YOLO11Seg(args.model)
  317.  
  318.     # 使用 OpenCV 读取图像
  319.     img = cv2.imread(args.source)
  320.  
  321.     # 模型推理
  322.     boxes, segments, _ = model(img, conf_threshold=args.conf, iou_threshold=args.iou)
  323.  
  324.     # 如果检测到目标，绘制边界框和分割区域
  325.     if len(boxes) > 0:
  326.         model.draw_and_visualize(img, boxes, segments, vis=False, save=True)

需要修改类别映射关系，以及类别对应的颜色

比如，定义两个类别（car、person）， 示例代码：

  1. # 类外定义类别映射关系，使用字典格式
  2. CLASS_NAMES = {
  3.     0: 'car',  # 类别 0 名称
  4.     1: 'person'   # 类别 1 名称
  5.                          # 可以添加更多类别...
  6. }

类别对应的两种颜色， 示例代码：

  1. # 定义类别对应的颜色，格式为 (R, G, B)
  2. CLASS_COLORS = {
  3.     0: (255, 255, 0),   # 类别 0 的颜色为青黄色
  4.     1: (255, 0, 0)      # 类别 1 的颜色为红色
  5.                         # 可以为其他类别指定颜色...
  6. }

运行代码，打印信息：

YOLO11 ? 实例分割 ONNXRuntime
模型名称： weights/seg_best_exp6.onnx
模型输入尺寸：宽度 = 640, 高度 = 640

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

 YOLO11相关文章推荐：

一篇文章快速认识YOLO11 | 关键改进点 | 安装使用 | 模型训练和推理-CSDN博客

一篇文章快速认识 YOLO11 | 实例分割 | 模型训练 | 自定义数据集-CSDN博客

YOLO11模型推理 | 目标检测与跟踪 | 实例分割 | 关键点估计 | OBB旋转目标检测-CSDN博客

YOLO11模型训练 | 目标检测与跟踪 | 实例分割 | 关键点姿态估计-CSDN博客

YOLO11 实例分割 | 自动标注 | 预标注 | 标签格式转换 | 手动校正标签-CSDN博客

分享完成，欢迎大家多多点赞和收藏，谢谢~

注：本文转载自blog.csdn.net的一颗小树x的文章"https://blog.csdn.net/qq_41204464/article/details/142922223"。版权归原作者所有，此博客不拥有其著作权，亦不承担相应法律责任。如有侵权，请联系我们删除。
复制链接
相关推荐
发表评论

评论记录：

未查询到任何数据！