使用Supervision库处理YOLOv11s-seg ONNX模型的输出

背景介绍

在计算机视觉领域，YOLO系列模型因其高效的检测性能而广受欢迎。YOLOv11s-seg是YOLO系列中支持实例分割的版本，能够同时输出目标检测框和分割掩码。本文将详细介绍如何利用Supervision库处理YOLOv11s-seg ONNX模型的输出结果。

模型导出与输出结构

YOLOv11s-seg模型可以通过Ultralytics框架导出为ONNX格式。导出时需要注意几个关键参数：

model.export(
    format="onnx",
    nms=True,
    data="data.yaml"
)

导出的ONNX模型有两个输出：

1. 检测结果：形状为(1, 300, 38)

   • 前4个值：边界框坐标(x1,y1,x2,y2)
   • 第5个值：置信度分数
   • 第6个值：类别ID
   • 第7-38个值：32个分割掩码系数

2. 掩码原型：形状为(1, 38, 160, 160)

输出处理实现

初始化模型

首先需要创建一个类来加载ONNX模型并处理输入输出：

class YOLOv11Processor:
    def __init__(self, model_path, conf_thres=0.7, iou_thres=0.5):
        self.conf_threshold = conf_thres
        self.session = onnxruntime.InferenceSession(model_path)
        self._setup_io_details()

输入预处理

输入图像需要调整为模型期望的尺寸并归一化：

def prepare_input(self, image):
    self.orig_height, self.orig_width = image.shape[:2]
    input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    input_img = cv2.resize(input_img, (self.input_width, self.input_height))
    input_img = input_img / 255.0
    input_img = input_img.transpose(2, 0, 1)
    return input_img[np.newaxis, :, :, :].astype(np.float32)

输出后处理

处理模型输出是核心部分，包括：

1. 过滤低置信度检测结果
2. 提取边界框并调整到原始图像尺寸
3. 生成分割掩码

def process_output(self, outputs):
    predictions = np.squeeze(outputs[0], axis=0)
    mask_protos = outputs[1]
    
    # 过滤低置信度结果
    conf_scores = predictions[:, 4]
    valid = conf_scores > self.conf_threshold
    predictions = predictions[valid]
    
    if len(predictions) == 0:
        return [], [], [], []
    
    # 提取边界框
    boxes = predictions[:, :4]
    boxes = self._rescale_boxes(boxes)
    
    # 提取类别ID
    class_ids = predictions[:, 5].astype(np.int32)
    
    # 生成分割掩码
    masks = self._extract_masks(predictions, mask_protos)
    
    return boxes, conf_scores[valid], class_ids, masks

掩码生成细节

掩码生成是通过32个掩码系数与32个掩码原型的加权求和实现的：

def _extract_masks(self, predictions, mask_protos):
    seg_coeffs = predictions[:, 6:38]
    mask_protos = mask_protos[0, :32, :, :]
    
    # 计算加权和
    masks = np.einsum('nc,chw->nhw', seg_coeffs, mask_protos)
    
    # Sigmoid激活
    masks = 1 / (1 + np.exp(-masks))
    
    # 二值化
    masks = masks > 0.5
    
    # 调整到原始图像尺寸
    final_masks = []
    for mask in masks:
        resized_mask = cv2.resize(
            mask.astype(np.uint8) * 255,
            (self.orig_width, self.orig_height),
            interpolation=cv2.INTER_NEAREST
        )
        final_masks.append(resized_mask)
    
    return np.array(final_masks)

使用Supervision可视化结果

处理完模型输出后，可以使用Supervision库进行可视化：

# 创建各种标注器
mask_annotator = sv.MaskAnnotator(color=sv.Color.GREEN)
box_annotator = sv.BoxAnnotator(thickness=2)
label_annotator = sv.LabelAnnotator(text_scale=0.7, text_thickness=2)

# 创建检测结果对象
detections = sv.Detections(
    xyxy=boxes,
    confidence=scores,
    class_id=class_ids,
    mask=masks
)

# 按顺序标注：掩码→边界框→标签
annotated_image = image.copy()
annotated_image = mask_annotator.annotate(annotated_image, detections)
annotated_image = box_annotator.annotate(annotated_image, detections)
annotated_image = label_annotator.annotate(annotated_image, detections)

常见问题与解决方案

1. 边界框位置不正确：

   • 确保正确地从模型输出中提取坐标
   • 验证坐标缩放是否正确应用到原始图像尺寸

2. 掩码显示异常：

   • 检查掩码系数是否正确提取
   • 确认掩码原型的使用是否正确
   • 验证掩码调整大小时的插值方法

3. 性能优化：

   • 批量处理图像可以提高效率
   • 考虑使用GPU加速ONNX推理

总结

本文详细介绍了如何使用Supervision库处理YOLOv11s-seg ONNX模型的输出。关键点包括正确解析模型输出结构、处理分割掩码以及使用Supervision进行可视化。通过这种方法，开发者可以轻松地将YOLOv11s-seg模型集成到自己的应用中，并实现高质量的实例分割结果可视化。

对于更复杂的应用场景，可以考虑进一步优化处理流程，例如实现异步处理、添加跟踪功能或集成更多后处理算法。