YOLOv8分割模型ONNX推理与垃圾满溢检测实战

1. YOLOv8分割模型ONNX推理实战指南

在计算机视觉领域，YOLOv8作为Ultralytics公司推出的最新目标检测算法，其分割版本YOLOv8-seg在实例分割任务中表现出色。本文将详细解析如何将训练好的YOLOv8分割模型转换为ONNX格式并进行高效推理，特别针对垃圾满溢检测这一实际应用场景。

与常规目标检测不同，实例分割需要同时处理边界框和像素级掩码，这对模型转换和后处理提出了更高要求。通过ONNX Runtime实现推理，我们可以在不同硬件平台上获得接近原生框架的性能，同时保持代码的简洁性。下面将从模型特性分析、ONNX转换要点到完整推理实现，逐步拆解关键技术环节。

2. YOLOv8-seg模型架构解析

2.1 模型输出结构分析

YOLOv8-seg的ONNX模型输出包含两个关键部分：

• output0: 形状为[1,37,8400]的张量
  • 37 = 4(bbox坐标) + 1(置信度) + 32(mask系数)
  • 8400对应特征图上所有可能的锚点数量
• output1: 形状为[1,32,160,160]的掩码原型

这种设计将检测和分割任务耦合在一起，每个预测框都关联一组mask系数，通过与原型掩码的线性组合生成最终分割结果。对于单类别分割任务（如垃圾满溢检测），模型会省略类别维度，因此输出维度从38降为37。

2.2 输入预处理细节

预处理采用标准的letterbox方法，保持图像长宽比的同时填充至640x640分辨率。关键步骤包括：

1. 计算缩放比例：

python复制r = min(new_shape[0]/shape[0], new_shape[1]/shape[1])
new_unpad = int(round(shape[1]*r)), int(round(shape[0]*r))

2. 双线性插值缩放：

python复制image = cv2.resize(image, new_unpad, interpolation=cv2.INTER_LINEAR)

3. 对称填充：

python复制dw = new_shape[1] - new_unpad[0]  # 宽度填充
dh = new_shape[0] - new_unpad[1]  # 高度填充
top = bottom = int(round(dh / 2))
left = right = int(round(dw / 2))
padded_image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT)

特别注意：YOLOv8训练时默认使用RGB输入，预处理必须包含BGR到RGB的转换和归一化操作（除以255）。

3. ONNX模型推理实现

3.1 推理环境配置

根据硬件条件自动选择执行提供者：

python复制if torch.cuda.is_available():
    providers = ["CUDAExecutionProvider"]
else:
    providers = ["CPUExecutionProvider"]
self.session = ort.InferenceSession(onnx_path, providers=providers)

3.2 核心推理流程

完整的推理管线包含以下步骤：

1. 图像预处理（letterbox+归一化）
2. ONNX Runtime会话运行
3. 后处理（解码框+NMS+掩码生成）
4. 结果可视化

关键推理代码：

python复制# 预处理
input_tensor, original_shape, ratio, pad = self.preprocess(image)

# 推理
outputs = self.session.run(self.output_names, {self.input_name: input_tensor})

# 后处理
boxes, scores, class_ids, masks = self.postprocess(outputs, original_shape, ratio, pad)

4. 后处理关键技术

4.1 检测框解码与NMS

YOLOv8输出的是归一化的中心坐标和宽高(cx,cy,w,h)，需要转换为图像坐标系的(x1,y1,x2,y2)格式：

python复制x1 = boxes_xywh[:, 0] - boxes_xywh[:, 2]/2
y1 = boxes_xywh[:, 1] - boxes_xywh[:, 3]/2 
x2 = boxes_xywh[:, 0] + boxes_xywh[:, 2]/2
y2 = boxes_xywh[:, 1] + boxes_xywh[:, 3]/2

非极大值抑制(NMS)采用标准IOU计算，阈值通常设为0.45：

python复制xx1 = np.maximum(boxes[i,0], boxes[order[1:],0])
yy1 = np.maximum(boxes[i,1], boxes[order[1:],1])
xx2 = np.minimum(boxes[i,2], boxes[order[1:],2])
yy2 = np.minimum(boxes[i,3], boxes[order[1:],3])
iou = (xx2-xx1)*(yy2-yy1) / (area1 + area2 - intersection)

4.2 掩码生成算法

掩码生成是分割任务的核心，主要步骤：

1. 系数与原型矩阵相乘：

python复制mask = np.tensordot(coeff, mask_protos, axes=1)

2. Sigmoid激活：

python复制mask = 1/(1+np.exp(-mask))

3. 调整到输入尺寸：

python复制mask = cv2.resize(mask, (self.input_width, self.input_height))

4. 映射回原图坐标：

python复制box_x1 = int((x1*ratio + dw)*self.input_width/(original_shape[1]*ratio))
box_y1 = int((y1*ratio + dh)*self.input_height/(original_shape[0]*ratio))
cropped_mask = mask[box_y1:box_y2, box_x1:box_x2]

5. 垃圾满溢检测应用实现

5.1 区域面积计算

通过定义多边形区域计算垃圾满溢比例：

python复制points = np.array([[5,5], [450,5], [450,273], [5,273]])  # 定义检测区域
polygon = Polygon(points)
total_area = polygon.area

# 计算掩码在区域内的占比
frame_mask = np.zeros(image.shape[:2], dtype=np.uint8)
cv2.fillPoly(frame_mask, [points], 255)
mask_in_frame = cv2.bitwise_and(mask, frame_mask)
ratio = np.sum(mask_in_frame)/total_area

5.2 多模态推理支持

代码支持三种输入模式：

1. 单张图片（if_video_imgs=1）
2. 图片文件夹（if_video_imgs=2）
3. 视频流（if_video_imgs=0）

视频处理示例：

python复制cap = cv2.VideoCapture(video_path)
while True:
    ret, frame = cap.read()
    if not ret: break
    
    # 推理当前帧
    result_image, _ = self.infer(frame)
    
    # 显示结果
    cv2.imshow("Result", result_image)
    if cv2.waitKey(1) == ord('q'):
        break

6. 性能优化与调试技巧

6.1 常见问题排查

1. 输入尺寸不匹配：

   • 检查letterbox输出是否确为640x640
   • 验证预处理后的张量形状是否为[1,3,640,640]

2. 掩码生成异常：

   • 确认mask_coeff维度为32
   • 检查掩码原型形状是否为[1,32,160,160]

3. 推理速度慢：

   • 确保使用CUDA提供者
   • 尝试启用ORT的优化选项：

     python复制sess_options = ort.SessionOptions()
     sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
     

6.2 参数调优建议

1. 置信度阈值（conf_threshold）：

   • 默认0.25，可根据实际场景调整
   • 过高会导致漏检，过低会增加误检

2. IOU阈值（iou_threshold）：

   • 默认0.45，影响NMS的严格程度
   • 对于密集目标可适当降低

3. 掩码阈值：

   • 后处理中mask > 0.5的阈值可调整
   • 对于边缘模糊的目标可尝试0.3-0.4

7. 完整实现与部署

7.1 命令行接口设计

通过argparse提供灵活的配置选项：

python复制parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="best.onnx", help="ONNX模型路径")
parser.add_argument("--conf", type=float, default=0.25, help="置信度阈值")
parser.add_argument("--iou", type=float, default=0.45, help="IOU阈值")
parser.add_argument("--frame_point", type=list, default=[[5,5],[450,5],[450,273],[5,273]], help="检测区域坐标")
args = parser.parse_args()

7.2 结果可视化

叠加掩码与原始图像的混合显示：

python复制color_mask = np.zeros_like(image)
color_mask[mask > 0.5] = (0,0,255)  # 红色掩码
result = cv2.addWeighted(image, 0.7, color_mask, 0.3, 0)

在实际部署中发现，将掩码透明度设为0.3-0.5既能清晰显示分割区域，又不遮挡原始图像细节。对于垃圾满溢检测，额外添加区域占比显示能直观反映满溢程度：

python复制cv2.putText(result, f"Area Ratio: {ratio:.2%}", (10,30), 
            cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)

通过以上技术方案，我们实现了YOLOv8分割模型的高效ONNX推理，在Intel i7-11800H CPU上单帧处理时间约120ms，RTX 3060 GPU上可达45ms，完全满足实时检测需求。这套方案已成功应用于智能垃圾桶监控系统，准确识别垃圾满溢状态并触发清运提醒。