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模型(简洁版)

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

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

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

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

YOLO11导出ONNX模型(专业版)

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

对于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.torchscriptimgszoptimizebatch
ONNXonnxyolo11n.onnximgszhalfdynamicsimplifyopsetbatch
OpenVINOopenvinoyolo11n_openvino_model/imgszhalfint8batch
TensorRTengineyolo11n.engineimgszhalfdynamicsimplifyworkspaceint8batch
CoreMLcoremlyolo11n.mlpackageimgszhalfint8nmsbatch
TF SavedModelsaved_modelyolo11n_saved_model/imgszkerasint8batch
TF GraphDefpbyolo11n.pbimgszbatch
TF Litetfliteyolo11n.tfliteimgszhalfint8batch
TF Edge TPUedgetpuyolo11n_edgetpu.tfliteimgsz
TF.jstfjsyolo11n_web_model/imgszhalfint8batch
PaddlePaddlepaddleyolo11n_paddle_model/imgszbatch
NCNNncnnyolo11n_ncnn_model/imgszhalfbatch

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

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

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

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

2.1、整体ONNX推理流程代码

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

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

class YOLO11Seg:def __init__(self, onnx_model):# 创建 Ort 推理会话,选择 CPU 或 GPU 提供者self.session = ort.InferenceSession(onnx_model,providers=["CUDAExecutionProvider", "CPUExecutionProvider"]if ort.get_device() == "GPU"else ["CPUExecutionProvider"],)# 根据 ONNX 模型类型选择 Numpy 数据类型(支持 FP32 和 FP16)self.ndtype = np.half if self.session.get_inputs()[0].type == "tensor(float16)" else np.single# 获取模型的输入宽度和高度(YOLO11-seg 只有一个输入)self.model_height, self.model_width = [x.shape for x in self.session.get_inputs()][0][-2:]# 打印模型的输入尺寸print("YOLO11 🚀 实例分割 ONNXRuntime")print("模型名称:", onnx_model)print(f"模型输入尺寸:宽度 = {self.model_width}, 高度 = {self.model_height}")# 加载类别名称self.classes = CLASS_NAMES# 加载类别对应的颜色self.class_colors = CLASS_COLORSdef __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45, nm=32):"""完整的推理流程:预处理 -> 推理 -> 后处理Args:im0 (Numpy.ndarray): 原始输入图像conf_threshold (float): 置信度阈值iou_threshold (float): NMS 中的 IoU 阈值nm (int): 掩膜数量Returns:boxes (List): 边界框列表segments (List): 分割区域列表masks (np.ndarray): [N, H, W] 输出掩膜"""# 图像预处理im, ratio, (pad_w, pad_h) = self.preprocess(im0)# ONNX 推理preds = self.session.run(None, {self.session.get_inputs()[0].name: im})# 后处理boxes, segments, masks = self.postprocess(preds,im0=im0,ratio=ratio,pad_w=pad_w,pad_h=pad_h,conf_threshold=conf_threshold,iou_threshold=iou_threshold,nm=nm,)return boxes, segments, masks

2.1、 图像预处理函数

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

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

2.3、 后处理函数

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

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

1、后处理框架 (postprocess)

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

2、掩膜处理

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

3、分割区域转换 (masks2segments)

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

4、掩膜缩放与裁剪

  • 通过 scale_maskcrop_mask,将掩膜调整到与原始图像匹配的大小,并裁剪到边界框内部,生成最终的分割掩膜。
    def postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold, nm=32):"""推理后的结果后处理Args:preds (Numpy.ndarray): 来自 ONNX 的推理结果im0 (Numpy.ndarray): [h, w, c] 原始输入图像ratio (tuple): 宽高比例pad_w (float): 宽度的填充pad_h (float): 高度的填充conf_threshold (float): 置信度阈值iou_threshold (float): IoU 阈值nm (int): 掩膜数量Returns:boxes (List): 边界框列表segments (List): 分割区域列表masks (np.ndarray): 掩膜数组"""x, protos = preds[0], preds[1]  # 获取模型的两个输出:预测和原型# 转换维度x = np.einsum("bcn->bnc", x)# 置信度过滤x = x[np.amax(x[..., 4:-nm], axis=-1) > conf_threshold]# 合并边界框、置信度、类别和掩膜x = np.c_[x[..., :4], np.amax(x[..., 4:-nm], axis=-1), np.argmax(x[..., 4:-nm], axis=-1), x[..., -nm:]]# NMS 过滤x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]# 解析并返回结果if len(x) > 0:# 边界框格式转换:从 cxcywh -> xyxyx[..., [0, 1]] -= x[..., [2, 3]] / 2x[..., [2, 3]] += x[..., [0, 1]]# 缩放边界框,使其与原始图像尺寸匹配x[..., :4] -= [pad_w, pad_h, pad_w, pad_h]x[..., :4] /= min(ratio)# 限制边界框在图像边界内x[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])# 处理掩膜masks = self.process_mask(protos[0], x[:, 6:], x[:, :4], im0.shape)# 将掩膜转换为分割区域segments = self.masks2segments(masks)return x[..., :6], segments, masks  # 返回边界框、分割区域和掩膜else:return [], [], []@staticmethoddef masks2segments(masks):"""将掩膜转换为分割区域Args:masks (numpy.ndarray): 模型输出的掩膜,形状为 (n, h, w)Returns:segments (List): 分割区域的列表"""segments = []for x in masks.astype("uint8"):c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[0]  # 找到轮廓if c:c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)else:c = np.zeros((0, 2))  # 如果没有找到分割区域,返回空数组segments.append(c.astype("float32"))return segments@staticmethoddef crop_mask(masks, boxes):"""裁剪掩膜,使其与边界框对齐Args:masks (Numpy.ndarray): [n, h, w] 掩膜数组boxes (Numpy.ndarray): [n, 4] 边界框Returns:(Numpy.ndarray): 裁剪后的掩膜"""n, h, w = masks.shapex1, y1, x2, y2 = np.split(boxes[:, :, None], 4, 1)r = np.arange(w, dtype=x1.dtype)[None, None, :]c = np.arange(h, dtype=x1.dtype)[None, :, None]return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))def process_mask(self, protos, masks_in, bboxes, im0_shape):"""处理模型输出的掩膜Args:protos (numpy.ndarray): [mask_dim, mask_h, mask_w] 掩膜原型masks_in (numpy.ndarray): [n, mask_dim] 掩膜数量bboxes (numpy.ndarray): 缩放到原始图像尺寸的边界框im0_shape (tuple): 原始输入图像的尺寸 (h,w,c)Returns:(numpy.ndarray): 处理后的掩膜"""c, mh, mw = protos.shapemasks = np.matmul(masks_in, protos.reshape((c, -1))).reshape((-1, mh, mw)).transpose(1, 2, 0)  # HWNmasks = np.ascontiguousarray(masks)masks = self.scale_mask(masks, im0_shape)  # 将掩膜从 P3 尺寸缩放到原始输入图像大小masks = np.einsum("HWN -> NHW", masks)  # HWN -> NHWmasks = self.crop_mask(masks, bboxes)  # 裁剪掩膜return np.greater(masks, 0.5)  # 返回二值化后的掩膜@staticmethoddef scale_mask(masks, im0_shape, ratio_pad=None):"""将掩膜缩放至原始图像大小Args:masks (np.ndarray): 缩放和填充后的掩膜im0_shape (tuple): 原始图像大小ratio_pad (tuple): 填充与原始图像的比例Returns:masks (np.ndarray): 缩放后的掩膜"""im1_shape = masks.shape[:2]if ratio_pad is None:  # 计算比例gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1])  # 比例pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2  # 填充else:pad = ratio_pad[1]# 计算掩膜的边界top, left = int(round(pad[1] - 0.1)), int(round(pad[0] - 0.1))  # y, xbottom, right = int(round(im1_shape[0] - pad[1] + 0.1)), int(round(im1_shape[1] - pad[0] + 0.1))if len(masks.shape) < 2:raise ValueError(f'"len of masks shape" 应该是 2 或 3,但得到 {len(masks.shape)}')masks = masks[top:bottom, left:right]masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]), interpolation=cv2.INTER_LINEAR)  # 使用 INTER_LINEAR 插值调整大小if len(masks.shape) == 2:masks = masks[:, :, None]return masks

2.4、分割效果可视化函数

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

    def draw_and_visualize(self, im, bboxes, segments, vis=False, save=True):"""绘制和可视化结果Args:im (np.ndarray): 原始图像,形状为 [h, w, c]bboxes (numpy.ndarray): [n, 4],n 是边界框数量segments (List): 分割区域的列表vis (bool): 是否使用 OpenCV 显示图像save (bool): 是否保存带注释的图像Returns:None"""# 创建图像副本im_canvas = im.copy()for (*box, conf, cls_), segment in zip(bboxes, segments):# 获取类别对应的颜色color = self.get_color_for_class(int(cls_))# 绘制轮廓和填充掩膜# cv2.polylines(im, np.int32([segment]), True, (255, 255, 255), 2)  # 绘制白色边框cv2.fillPoly(im_canvas, np.int32([segment]), color)  # 使用类别对应的颜色填充多边形# 绘制边界框cv2.rectangle(im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 1, cv2.LINE_AA)# 在图像上绘制类别名称和置信度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)# 将图像和绘制的多边形混合im = cv2.addWeighted(im_canvas, 0.3, im, 0.7, 0)# 显示图像if vis:cv2.imshow("seg_result_picture", im)cv2.waitKey(0)cv2.destroyAllWindows()# 保存图像if save:cv2.imwrite("seg_result_picture.jpg", im)

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

完整代码,如下所示:

# Ultralytics YOLO 🚀, AGPL-3.0 license
""" 
YOLO11 分割模型 ONNXRuntime 功能1: 支持不用尺寸图像的输入功能2: 支持可视化分割结果
"""import argparse
import cv2
import numpy as np
import onnxruntime as ort# 类外定义类别映射关系,使用字典格式
CLASS_NAMES = {0: 'class_name1',  # 类别 0 名称1: 'class_name2'   # 类别 1 名称# 可以添加更多类别...
}# 定义类别对应的颜色,格式为 (R, G, B)
CLASS_COLORS = {0: (255, 255, 0),   # 类别 0 的颜色为青黄色1: (255, 0, 0)      # 类别 1 的颜色为红色# 可以为其他类别指定颜色...
}class YOLO11Seg:def __init__(self, onnx_model):# 创建 Ort 推理会话,选择 CPU 或 GPU 提供者self.session = ort.InferenceSession(onnx_model,providers=["CUDAExecutionProvider", "CPUExecutionProvider"]if ort.get_device() == "GPU"else ["CPUExecutionProvider"],)# 根据 ONNX 模型类型选择 Numpy 数据类型(支持 FP32 和 FP16)self.ndtype = np.half if self.session.get_inputs()[0].type == "tensor(float16)" else np.single# 获取模型的输入宽度和高度(YOLO11-seg 只有一个输入)self.model_height, self.model_width = [x.shape for x in self.session.get_inputs()][0][-2:]# 打印模型的输入尺寸print("YOLO11 🚀 实例分割 ONNXRuntime")print("模型名称:", onnx_model)print(f"模型输入尺寸:宽度 = {self.model_width}, 高度 = {self.model_height}")# 加载类别名称self.classes = CLASS_NAMES# 加载类别对应的颜色self.class_colors = CLASS_COLORSdef get_color_for_class(self, class_id):return self.class_colors.get(class_id, (255, 255, 255))  # 如果没有找到类别颜色,返回白色def __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45, nm=32):"""完整的推理流程:预处理 -> 推理 -> 后处理Args:im0 (Numpy.ndarray): 原始输入图像conf_threshold (float): 置信度阈值iou_threshold (float): NMS 中的 IoU 阈值nm (int): 掩膜数量Returns:boxes (List): 边界框列表segments (List): 分割区域列表masks (np.ndarray): [N, H, W] 输出掩膜"""# 图像预处理im, ratio, (pad_w, pad_h) = self.preprocess(im0)# ONNX 推理preds = self.session.run(None, {self.session.get_inputs()[0].name: im})# 后处理boxes, segments, masks = self.postprocess(preds,im0=im0,ratio=ratio,pad_w=pad_w,pad_h=pad_h,conf_threshold=conf_threshold,iou_threshold=iou_threshold,nm=nm,)return boxes, segments, masksdef preprocess(self, img):"""图像预处理Args:img (Numpy.ndarray): 输入图像Returns:img_process (Numpy.ndarray): 处理后的图像ratio (tuple): 宽高比例pad_w (float): 宽度的填充pad_h (float): 高度的填充"""# 调整输入图像大小并使用 letterbox 填充shape = img.shape[:2]  # 原始图像大小new_shape = (self.model_height, self.model_width)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])ratio = r, rnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # 填充宽高if shape[::-1] != new_unpad:  # 调整图像大小img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))# 转换:HWC -> CHW -> BGR 转 RGB -> 除以 255 -> contiguous -> 添加维度img = np.ascontiguousarray(np.einsum("HWC->CHW", img)[::-1], dtype=self.ndtype) / 255.0img_process = img[None] if len(img.shape) == 3 else imgreturn img_process, ratio, (pad_w, pad_h)def postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold, nm=32):"""推理后的结果后处理Args:preds (Numpy.ndarray): 来自 ONNX 的推理结果im0 (Numpy.ndarray): [h, w, c] 原始输入图像ratio (tuple): 宽高比例pad_w (float): 宽度的填充pad_h (float): 高度的填充conf_threshold (float): 置信度阈值iou_threshold (float): IoU 阈值nm (int): 掩膜数量Returns:boxes (List): 边界框列表segments (List): 分割区域列表masks (np.ndarray): 掩膜数组"""x, protos = preds[0], preds[1]  # 获取模型的两个输出:预测和原型# 转换维度x = np.einsum("bcn->bnc", x)# 置信度过滤x = x[np.amax(x[..., 4:-nm], axis=-1) > conf_threshold]# 合并边界框、置信度、类别和掩膜x = np.c_[x[..., :4], np.amax(x[..., 4:-nm], axis=-1), np.argmax(x[..., 4:-nm], axis=-1), x[..., -nm:]]# NMS 过滤x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]# 解析并返回结果if len(x) > 0:# 边界框格式转换:从 cxcywh -> xyxyx[..., [0, 1]] -= x[..., [2, 3]] / 2x[..., [2, 3]] += x[..., [0, 1]]# 缩放边界框,使其与原始图像尺寸匹配x[..., :4] -= [pad_w, pad_h, pad_w, pad_h]x[..., :4] /= min(ratio)# 限制边界框在图像边界内x[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])# 处理掩膜masks = self.process_mask(protos[0], x[:, 6:], x[:, :4], im0.shape)# 将掩膜转换为分割区域segments = self.masks2segments(masks)return x[..., :6], segments, masks  # 返回边界框、分割区域和掩膜else:return [], [], []@staticmethoddef masks2segments(masks):"""将掩膜转换为分割区域Args:masks (numpy.ndarray): 模型输出的掩膜,形状为 (n, h, w)Returns:segments (List): 分割区域的列表"""segments = []for x in masks.astype("uint8"):c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[0]  # 找到轮廓if c:c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)else:c = np.zeros((0, 2))  # 如果没有找到分割区域,返回空数组segments.append(c.astype("float32"))return segments@staticmethoddef crop_mask(masks, boxes):"""裁剪掩膜,使其与边界框对齐Args:masks (Numpy.ndarray): [n, h, w] 掩膜数组boxes (Numpy.ndarray): [n, 4] 边界框Returns:(Numpy.ndarray): 裁剪后的掩膜"""n, h, w = masks.shapex1, y1, x2, y2 = np.split(boxes[:, :, None], 4, 1)r = np.arange(w, dtype=x1.dtype)[None, None, :]c = np.arange(h, dtype=x1.dtype)[None, :, None]return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))def process_mask(self, protos, masks_in, bboxes, im0_shape):"""处理模型输出的掩膜Args:protos (numpy.ndarray): [mask_dim, mask_h, mask_w] 掩膜原型masks_in (numpy.ndarray): [n, mask_dim] 掩膜数量bboxes (numpy.ndarray): 缩放到原始图像尺寸的边界框im0_shape (tuple): 原始输入图像的尺寸 (h,w,c)Returns:(numpy.ndarray): 处理后的掩膜"""c, mh, mw = protos.shapemasks = np.matmul(masks_in, protos.reshape((c, -1))).reshape((-1, mh, mw)).transpose(1, 2, 0)  # HWNmasks = np.ascontiguousarray(masks)masks = self.scale_mask(masks, im0_shape)  # 将掩膜从 P3 尺寸缩放到原始输入图像大小masks = np.einsum("HWN -> NHW", masks)  # HWN -> NHWmasks = self.crop_mask(masks, bboxes)  # 裁剪掩膜return np.greater(masks, 0.5)  # 返回二值化后的掩膜@staticmethoddef scale_mask(masks, im0_shape, ratio_pad=None):"""将掩膜缩放至原始图像大小Args:masks (np.ndarray): 缩放和填充后的掩膜im0_shape (tuple): 原始图像大小ratio_pad (tuple): 填充与原始图像的比例Returns:masks (np.ndarray): 缩放后的掩膜"""im1_shape = masks.shape[:2]if ratio_pad is None:  # 计算比例gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1])  # 比例pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2  # 填充else:pad = ratio_pad[1]# 计算掩膜的边界top, left = int(round(pad[1] - 0.1)), int(round(pad[0] - 0.1))  # y, xbottom, right = int(round(im1_shape[0] - pad[1] + 0.1)), int(round(im1_shape[1] - pad[0] + 0.1))if len(masks.shape) < 2:raise ValueError(f'"len of masks shape" 应该是 2 或 3,但得到 {len(masks.shape)}')masks = masks[top:bottom, left:right]masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]), interpolation=cv2.INTER_LINEAR)  # 使用 INTER_LINEAR 插值调整大小if len(masks.shape) == 2:masks = masks[:, :, None]return masksdef draw_and_visualize(self, im, bboxes, segments, vis=False, save=True):"""绘制和可视化结果Args:im (np.ndarray): 原始图像,形状为 [h, w, c]bboxes (numpy.ndarray): [n, 4],n 是边界框数量segments (List): 分割区域的列表vis (bool): 是否使用 OpenCV 显示图像save (bool): 是否保存带注释的图像Returns:None"""# 创建图像副本im_canvas = im.copy()for (*box, conf, cls_), segment in zip(bboxes, segments):# 获取类别对应的颜色color = self.get_color_for_class(int(cls_))# 绘制轮廓和填充掩膜# cv2.polylines(im, np.int32([segment]), True, (255, 255, 255), 2)  # 绘制白色边框cv2.fillPoly(im_canvas, np.int32([segment]), color)  # 使用类别对应的颜色填充多边形# 绘制边界框cv2.rectangle(im, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 1, cv2.LINE_AA)# 在图像上绘制类别名称和置信度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)# 将图像和绘制的多边形混合im = cv2.addWeighted(im_canvas, 0.3, im, 0.7, 0)# 显示图像if vis:cv2.imshow("seg_result_picture", im)cv2.waitKey(0)cv2.destroyAllWindows()# 保存图像if save:cv2.imwrite("seg_result_picture.jpg", im)if __name__ == "__main__":# 创建命令行参数解析器parser = argparse.ArgumentParser()parser.add_argument("--model", type=str, default=r"weights/seg_best_exp6.onnx" , help="ONNX 模型路径")parser.add_argument("--source", type=str, default=r"datasets/Mix_seg_point_offer_num53/images/point_offer_20240930_14_Image.png", help="输入图像路径")parser.add_argument("--conf", type=float, default=0.6, help="置信度阈值")parser.add_argument("--iou", type=float, default=0.45, help="NMS 的 IoU 阈值")args = parser.parse_args()# 加载模型model = YOLO11Seg(args.model)# 使用 OpenCV 读取图像img = cv2.imread(args.source)# 模型推理boxes, segments, _ = model(img, conf_threshold=args.conf, iou_threshold=args.iou)# 如果检测到目标,绘制边界框和分割区域if len(boxes) > 0:model.draw_and_visualize(img, boxes, segments, vis=False, save=True)

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

比如,定义两个类别(car、person), 示例代码:

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

类别对应的两种颜色, 示例代码:

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

运行代码,打印信息:

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博客

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

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.rhkb.cn/news/447598.html

如若内容造成侵权/违法违规/事实不符,请联系长河编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

HDMI Type B 连接器的引脚分配与接口尺寸

连接器引脚分配 B 型连接器引脚分配&#xff1a; 连接器接触顺序 连接器接口尺寸 B 型插座接口尺寸&#xff1a; B 型插头接口尺寸&#xff1a; B 型插座与插头配合状态&#xff1a; 微信公众号&#xff1a;

CreateEvent 事件对象可以用来在线程之间同步操作 信号事件

CreateEvent 是一个 Windows API 函数&#xff0c;用于创建或打开一个事件对象&#xff0c;事件对象可以用来在线程之间同步操作。该函数的签名如下&#xff1a; HANDLE CreateEvent(LPSECURITY_ATTRIBUTES lpEventAttributes, // 安全属性&#xff0c;可以为 NULLBOOL bManu…

使用XML实现MyBatis的基础操作

目录 前言 1.准备工作 1.1⽂件配置 1.2添加 mapper 接⼝ 2.增删改查操作 2.1增(Insert) 2.2删(Delete) 2.3改(Update) 2.4查(Select) 前言 接下来我们会使用的数据表如下&#xff1a; 对应的实体类为&#xff1a;UserInfo 所有的准备工作都在如下文章。 MyBatis 操作…

React国际化中英文切换实现

目录 概况 安装 文件结构 引入 使用 正常使用 传参使用 概况 react-intl-universal 是一个国际化库&#xff0c;专门为 React 应用提供多语言支持。与 React 原生的 react-intl 相比&#xff0c;react-intl-universal 支持从远程服务器加载语言包&#xff0c;动态切换语…

MySQL数据的导入

【图书推荐】《MySQL 9从入门到性能优化&#xff08;视频教学版&#xff09;》-CSDN博客 《MySQL 9从入门到性能优化&#xff08;视频教学版&#xff09;&#xff08;数据库技术丛书&#xff09;》(王英英)【摘要 书评 试读】- 京东图书 (jd.com) MySQL9数据库技术_夏天又到了…

简单使用webpack

下载依赖&#xff1a;npm install webpack webpack-cli --save-dev新建一个目录app。加入html&#xff1a; <html><head><meta charset"utf-8"></head><body><script type"text/javascript" src"./dist/main.js&quo…

05 django管理系统 - 部门管理 - 修改部门

04我们已经实现了新增部门的功能&#xff0c;下面开始修改部门模块的实现。 按道理来说&#xff0c;应该是做成弹框样式的&#xff0c;通过ajax悄咪咪的发数据&#xff0c;然后更新前端数据&#xff0c;但是考虑到实际情况&#xff0c;先用页面跳转的方式实现&#xff0c;后面…

c到c++衔接速成

温馨提示&#xff1a;本篇文章乃博主多次学习后的总结&#xff0c;其中一些专业名词的概念不会做介绍&#xff0c;本篇文章只解决实操问题 c到c中&#xff0c;语法上得到了许多升级&#xff0c;那么哪些地方升级了&#xff0c;我们来快速复习一下 目录 头文件的升级 命名空间…

Android基于gradle task检查各个module之间资源文件冲突情况

做组件化开发的时候&#xff0c;我们经常会遇到各个不同的module之间资源文件冲突的问题&#xff0c;运行也不报错&#xff0c;但是会出现覆盖的问题&#xff0c;导致运行之后发送错误的效果。 所以我们需要利用一个gradlke 脚本task&#xff0c;来自动化检查资源文件冲突。 …

利用 OBS 推送 WEBRTC 流到 smart rtmpd

webrtc whip 推流 & whep 拉流简介 RFC 定义 通用的 webrtc 对于 SDP 协议的交换已经有对应的 RFC 草案出炉了。这就是 WHIP( push stream ) & WHEP ( pull stream ) . WHIP RFC Link: https://www.ietf.org/archive/id/draft-ietf-wish-whip-01.html WHEP RFC Link:…

IRP读写函数

驱动代码&#xff1a; #include <ntddk.h>#define DEVICE_NAME L"\\device\\MyDricer1" //设备对象名称 #define LINK_NAME L"\\dosdevices\\Goose" //符号链接名称VOID UnDirver(PDRIVER_OBJECT pDriverObj) {UNICODE_STRING uLinkName RTL_CONST…

04_实现调度执行

什么是可调度&#xff1f;所谓可调度&#xff0c;指的是当 trigger 动作触发副作用函数重新执行的时候&#xff0c;有能力决定副作用函数执行的时机、次数以及方式 副作用函数执行方式 示例代码&#xff1a; const obj { a: 1 } const objProxy new Proxy(obj, /* ... */)…

02.顺序表、链表简述+对比

目录 一、线性表 二、顺序表 三、链表 四、顺序表和链表的区别 一、线性表 线性表&#xff08;linear list&#xff09;是n个具有相同特性的数据元素的有限序列&#xff08;相同特性指都为整型int、字符型char或其它类型&#xff09;。 线性表是一种在实际中广泛使用的数据…

雷池社区版配置遇到问题不要慌,查看本文解决

很多新人不太熟悉反向代理&#xff0c;所以导致配置站点出现问题 配置问题 记录常见的配置问题 配置后攻击测试没有拦截记录 检查访问请求有没有真实经过雷池 有很多新人配置站点后&#xff0c;真实的网站流量还是走的源站&#xff0c;导致雷池这边什么数据都没有 配置后…

软考《信息系统运行管理员》- 5.2 信息系统数据资源例行管理

5.2 信息系统数据资源例行管理 文章目录 5.2 信息系统数据资源例行管理数据资源例行管理计划数据资源载体的管理存储介质借用管理存储介质转储管理存储介质销毁管理 数据库例行维护健康检查数据库日志检查数据库一致性检查 数据库监测管理数据库备份与恢复数据库备份与恢复数据…

三子棋(C 语言)

目录 一、游戏设计的整体思路二、各个步骤的代码实现1. 菜单及循环选择的实现2. 棋盘的初始化和显示3. 轮流下棋及结果判断实现4. 结果判断实现 三、所有代码四、总结 一、游戏设计的整体思路 &#xff08;1&#xff09;提供一个菜单让玩家选择人机对战、玩家对战或者退出游戏…

生成式人工智能助长更复杂的网络攻击

网络安全公司 Keeper Security 表示&#xff0c;尽管黑客使用生成式人工智能使网络威胁更加频繁、更加复杂、更难以检测&#xff0c;但许多企业表示对此并未做好准备。 然而&#xff0c;据该公司表示&#xff0c;尽管各组织需要加强防御能力&#xff0c;但他们已经实施的许多基…

从调用NCCL到深入NCCL源码

本小白目前研究GPU多卡互连的方案&#xff0c;主要参考NCCL和RCCL进行学习&#xff0c;如有错误&#xff0c;请及时指正&#xff01; 内容还在整理中&#xff0c;近期不断更新&#xff01;&#xff01; 背景介绍 在大模型高性能计算时会需要用到多卡&#xff08;GPU&#xf…

接口性能测试,这个还真有用啊。

一、概述 性能测试按照不同视角&#xff0c;可以分为以下几类&#xff1a; a. 用户视角的性能 用户角度感受到的网站响应速度的快和慢。从用户在浏览器输入网址/打开应用&#xff0c;到整个页面呈现给用户的耗时。包含了用户端发送请求&#xff0c;服务端收到并执行请求&…

Python和MATLAB及C++和Fortran胶体粒子数学材料学显微镜学微观流变学及光学计算

&#x1f3af;要点 二维成像拥挤胶体粒子检测算法粒子的局部结构和动力学分析椭圆粒子成链动态过程定量分析算法小颗粒的光散射和吸收活跃物质模拟群体行为提取粒子轨迹粘弹性&#xff0c;计算剪切模量计算悬浮液球形粒子多体流体动力学概率规划全息图跟踪和表征粒子位置、大小…