yolov8分割任务的推理和后处理解析-EW帮帮网

在这里插入图片描述

文章目录

一、前言
二、分割模型的前向推理
三、后处理

一、前言

这篇文章主要分享yolov8模型用于图像分割时，模型输出和后处理。彻底理了下，可以总结为以下3点：

分割继承检测，前向推理时也会调用检测的方法把目标框检测出来；
但是前向推理分割和检测是各自进行的，训练也是分别去计算loss；
在后处理时为了提精度，在有目标处才去分割，然后为了提速掩膜才去系数和乘以原始掩膜的方法，系数和原始掩膜都是分割模型的前向推理输出；

yolov8官方代码路径：https://github.com/ultralytics/ultralytics

二、分割模型的前向推理

代码位置：yolo/ultralytics/nn/modules/head.py
在这里插入图片描述
解释：

继承关系：

Segment继承了Detect的所有基础功能，包括目标检测的能力
它扩展了Detect的功能，增加了实例分割的能力

主要组件：

def __init__(self, nc=80, nm=32, npr=256, ch=()):
    super().__init__(nc, ch)
    self.nm = nm  # 掩码数量
    self.npr = npr  # 原型数量
    self.proto = Proto(ch[0], self.npr, self.nm)  # 原型网络
    self.detect = Detect.forward  # 保留检测功能

推理输出：
从forward方法可以看出，Segment模型在推理时返回两个主要部分：

def forward(self, x):
    p = self.proto(x[0])  # 生成掩码原型
    bs = p.shape[0]  # batch size
    # 生成掩码系数
    mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)
    x = self.detect(self, x)  # 调用检测功能
    if self.training:
        return x, mc, p
    return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))

推理时返回的内容包括：

1. 检测结果：来自Detect类的输出

变量解释：

x：分别为3个head输出的特征图（大中小）
shape为：（bs, 4+类别数，特征图宽，特征图高）
y: 边界框坐标+类别预测(经过sigmoid)——纵向拼接
shape为：（bs, 4+类别数，框的个数）
训练模式，则输出x；
推理模式：
export为onnx时则输出：y
否则输出一个元组：(y, x)

2. 分割结果（最终）

变量解释：

掩码系数mc（mask coefficients）
shape为：（bs, 32(系数个数)，框的个数）
原型掩码p（prototype masks）
shape为：（bs, 32(系数个数)，mask图宽，mask图高）

训练模式，输出三个元素：x(detect的输出，对应x)，mc，p

推理模式：

export为onnx时输出元组包含2个元素：

第一个元素：纵向(第1维)拼接x(这里对应detect输出的y)和mc
shape为：（bs, 4+类别数+32(系数个数)，框的个数）
第二个元素：p

否则也是输出元组包含2个元素：

第一个元素只有1个元素：纵向拼接x0和mc
可以理解为：目标检测的结果+掩码系数
shape为（bs, 4+类别数+32，mask(或框)的个数）
–>（4后处理的输入[0]
第二个元素是个元组有3个元素：(x1, mc, p)
可以理解为：目标检测的head特征，掩码系数，原型掩码
–>（4后处理的输入[1]

3. 与Detect的主要区别

Detect只输出检测结果（边界框和类别）
Segment在Detect的基础上增加了分割能力，可以输出实例掩码
Segment使用原型网络（Proto）来生成掩码，这是分割特有的组件

4. 工作流程

首先通过原型网络生成基础掩码
同时进行目标检测
将检测结果和掩码系数结合，生成最终的实例分割结果
这种设计使得Segment模型能够同时完成目标检测和实例分割任务，是一个多任务的模型架构。

三、后处理

代码位置：yolo/ultralytics/yolo/v8/segment/predict.py
在这里插入图片描述
其中:
pred[0]实际上就是：纵向拼接x0和mc
pred[1]实际上就是：(x1, mc, p)

1. 非极大值抑制（NMS）过滤检测框

功能：
- 过滤掉低置信度（< conf）的检测框。
- 合并IoU超过阈值（iou）的重叠框。
- 若启用agnostic_nms，不同类别的框也会被合并（适用于类别无关任务）。
- 输出p为一个列表，每个元素对应一张图像的检测结果（形状：(num_boxes, 6 + num_masks)，其中num_boxes为保留的检测框数，6包含x1,y1,x2,y2,conf,cls,mask1系数,mask2系数,…32个系数）。
注意：
- preds[0]形状为(batch_size, 4 + num_classes + num_masks, num_boxes)
- num_masks 为 mask系数数量，通常是32个
- mask_coeffs用于和原型掩码线性组合生成实例分割
- 原型掩码是由模型预测出来的，对应output1

2. 分割原型（Mask Prototypes）提取

背景：
- preds[1]是分割头的输出，包含掩码原型。
- 若模型为PyTorch格式（未导出），preds[1]可能有3个元素（如不同尺度的原型），需取最后一个（最高分辨率）。
- 若模型已导出（如ONNX），preds[1]直接为原型张量。
形状：
- proto的典型形状为[batch, K, H, W]，其中：
  - K：原型数量（如32）。
  - H, W：原型的分辨率（如输入图像的1/4大小）

3. 掩码生成

在这里插入图片描述
注意：每一个框对应一组掩码系数。

分为两个模式：
有四个尺寸：特征图尺寸（框对应）；input尺寸；mask尺寸；原图尺寸

(1) 视网膜掩码：（标蓝是一个过程）
精度更高

放大box坐标到原图->生成mask(小图)->裁剪mask图(因为输入的时候为了保持图像不变形，会在两侧添加填充)->resize mask到原图->裁切mask对齐框(保留检测框内的区域，框外区域置为0)

(2) 普通掩码:（标蓝是一个过程）-- 推理默认
性能更好

生成mask(小图)->把坐标缩放到mask->裁切mask对齐框(保留检测框内的区域，框外区域置为0)->resize mask到input尺寸->把坐标放大到原图
注意：这个时候resize mask跟box坐标不在同一个尺寸标准下，画图的时候会把mask缩放到原图大小。

代码位置：yolo/ultralytics/yolo/utils/plotting.py
在这里插入图片描述

(3) 这两个模式都包括了两个操作：

(a) 缩放坐标

将检测框坐标从模型输入尺寸（img.shape[2:]）缩放到原始图像尺寸（orig_img.shape），处理填充（padding）和缩放比例，确保框位置正确映射。
拆切超出图像边缘部分的框。

(b) 生成掩码
代码位置：yolo/ultralytics/yolo/utils/ops.py

如果是视网膜掩码，则使用process_mask_native

输入参数：

def process_mask_native(protos, masks_in, bboxes, shape):
    """
    protos: 原型掩码 [mask_dim, mask_h, mask_w]
    masks_in: 预测的掩码系数 [n, mask_dim]
    bboxes: 检测框 [n, 4]
    shape: 原始图像尺寸 (h,w)
    """

掩码生成：

c, mh, mw = protos.shape  # 获取原型掩码的维度
# 将掩码系数与原型掩码相乘，得到最终掩码
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)

将原型掩码展平为2D矩阵
与掩码系数进行矩阵乘法
应用sigmoid激活函数
重塑为3D张量

计算缩放和填充：

# 计算缩放比例
gain = min(mh / shape[0], mw / shape[1])  # gain = 旧尺寸/新尺寸

# 计算填充值
pad = (mw - shape[1] * gain) / 2, (mh - shape[0] * gain) / 2  # wh padding
top, left = int(pad[1]), int(pad[0])  # y, x
bottom, right = int(mh - pad[1]), int(mw - pad[0])

计算保持宽高比的缩放比例
计算图像两侧的填充值
确定裁剪区域

裁剪掩码：

# 裁剪掉填充区域
masks = masks[:, top:bottom, left:right]

移除填充区域
保持有效区域

调整大小：

# 将掩码调整到原始图像大小
masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]

使用双线性插值
调整到原始图像尺寸
保持掩码质量

根据检测框裁剪(保留检测框内的区域，框外区域置为0)：

# 根据检测框裁剪掩码
masks = crop_mask(masks, bboxes)

将掩码裁剪到检测框区域
确保掩码与检测框对齐

二值化处理：

# 将掩码二值化
return masks.gt_(0.5)

使用0.5作为阈值
将掩码转换为二值图像

如果是普通掩码，则使用process_mask

生成mask
将检测框坐标从图像尺寸缩放到掩码尺寸
使用缩放后的检测框裁剪掩码，确保掩码与检测框对齐
mask上采样要原图

yolov8分割任务的推理和后处理解析