本文分享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( )函数中,各种参数说明:
format="onnx"
:指定导出模型的格式为 ONNX。imgsz=(640, 640)
:输入图像的尺寸设为 640x640。如果需要其他尺寸可以修改这个值。keras=False
:不导出为 Keras 格式的模型。optimize=False
:不应用 TorchScript 移动设备优化。half=False
:不启用 FP16(半精度)量化。int8=False
:不启用 INT8 量化。dynamic=False
:不启用动态输入尺寸。simplify=True
:简化模型以提升 ONNX 模型的性能。opset=None
:使用默认的 ONNX opset 版本,如果需要可以手动指定。workspace=4.0
:为 TensorRT 优化指定最大工作空间大小为 4 GiB。nms=False
:不为 CoreML 导出添加非极大值抑制(NMS)。batch=1
:设置批处理大小为 1。
参考官网文档:https://docs.ultralytics.com/modes/export/#arguments
当然了,YOLO11中不但支持ONNX模型,还支持下面表格中格式
支持的导出格式 | format 参数值 | 生成的模型示例 | model.export( )函数的参数 |
---|---|---|---|
PyTorch | - | yolo11n.pt | - |
TorchScript | torchscript | yolo11n.torchscript | imgsz , optimize , batch |
ONNX | onnx | yolo11n.onnx | imgsz , half , dynamic , simplify , opset , batch |
OpenVINO | openvino | yolo11n_openvino_model/ | imgsz , half , int8 , batch |
TensorRT | engine | yolo11n.engine | imgsz , half , dynamic , simplify , workspace , int8 , batch |
CoreML | coreml | yolo11n.mlpackage | imgsz , half , int8 , nms , batch |
TF SavedModel | saved_model | yolo11n_saved_model/ | imgsz , keras , int8 , batch |
TF GraphDef | pb | yolo11n.pb | imgsz , batch |
TF Lite | tflite | yolo11n.tflite | imgsz , half , int8 , batch |
TF Edge TPU | edgetpu | yolo11n_edgetpu.tflite | imgsz |
TF.js | tfjs | yolo11n_web_model/ | imgsz , half , int8 , batch |
PaddlePaddle | paddle | yolo11n_paddle_model/ | imgsz , batch |
NCNN | ncnn | yolo11n_ncnn_model/ | imgsz , half , batch |
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_mask
和crop_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博客
分享完成,欢迎大家多多点赞和收藏,谢谢~