文章目錄
- 一、前言
- 二、分割模型的前向推理
- 1. 檢測結果:來自Detect類的輸出
- 2. 分割結果(最終)
- 3. 與Detect的主要區別
- 4. 工作流程
- 三、后處理
- 1. 非極大值抑制(NMS)過濾檢測框
- 2. 分割原型(Mask Prototypes)提取
- 3. 掩碼生成
一、前言
這篇文章主要分享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, 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大小)
- proto的典型形狀為[batch, K, H, W],其中:
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上采樣要原圖