一、前言

這篇文章主要分享yolov8模型用于圖像分割時，模型輸出和后處理。徹底理了下，可以總結為以下3點：

1. 分割繼承檢測，前向推理時也會調用檢測的方法把目標框檢測出來；
2. 但是前向推理分割和檢測是各自進行的，訓練也是分別去計算loss；
3. 在后處理時為了提精度，在有目標處才去分割，然后為了提速掩膜才去系數和乘以原始掩膜的方法，系數和原始掩膜都是分割模型的前向推理輸出；

yolov8官方代碼路徑：https://github.com/ultralytics/ultralytics

二、分割模型的前向推理

代碼位置：yolo/ultralytics/nn/modules/head.py

解釋：

1. 繼承關系：

• Segment繼承了Detect的所有基礎功能，包括目標檢測的能力
• 它擴展了Detect的功能，增加了實例分割的能力

2. 主要組件：

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  # 保留檢測功能

3. 推理輸出：
   從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, preturn (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) 縮放坐標

1. 將檢測框坐標從模型輸入尺寸（img.shape[2:]）縮放到原始圖像尺寸（orig_img.shape），處理填充（padding）和縮放比例，確保框位置正確映射。
2. 拆切超出圖像邊緣部分的框。

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

如果是視網膜掩碼，則使用process_mask_native

1. 輸入參數：

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)"""

2. 掩碼生成：

c, mh, mw = protos.shape  # 獲取原型掩碼的維度
# 將掩碼系數與原型掩碼相乘，得到最終掩碼
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)

• 將原型掩碼展平為2D矩陣
• 與掩碼系數進行矩陣乘法
• 應用sigmoid激活函數
• 重塑為3D張量

3. 計算縮放和填充：

# 計算縮放比例
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])

• 計算保持寬高比的縮放比例
• 計算圖像兩側的填充值
• 確定裁剪區域

4. 裁剪掩碼：

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

• 移除填充區域
• 保持有效區域

5. 調整大小：

# 將掩碼調整到原始圖像大小
masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0]

• 使用雙線性插值
• 調整到原始圖像尺寸
• 保持掩碼質量

6. 根據檢測框裁剪(保留檢測框內的區域，框外區域置為0)：

# 根據檢測框裁剪掩碼
masks = crop_mask(masks, bboxes)

• 將掩碼裁剪到檢測框區域
• 確保掩碼與檢測框對齊

7. 二值化處理：

# 將掩碼二值化
return masks.gt_(0.5)

• 使用0.5作為閾值
• 將掩碼轉換為二值圖像

如果是普通掩碼，則使用process_mask

1. 生成mask
2. 將檢測框坐標從圖像尺寸縮放到掩碼尺寸
3. 使用縮放后的檢測框裁剪掩碼，確保掩碼與檢測框對齊
4. mask上采樣要原圖