作者：賀弘謙言臨在

導言

影像分割(Image Segmentation)是指對圖片進行畫素級的分類，根據分類粒度的不同可以分為語義分割(Semantic Segmentation)、例項分割(Instance Segmentation)、全景分割(Panoptic Segmentation)三類。影像分割是計算機視覺中的主要研究方向之一，在醫學影像分析、自動駕駛、影片監控、擴增實境、影像壓縮等領域有重要的應用價值。我們在EasyCV框架中對這三類分割SOTA演算法進行了整合，並提供了相關模型權重。透過EasyCV可以輕鬆預測影像的分割譜以及訓練定製化的分割模型。本文主要介紹如何使用EasyCV實現例項分割、全景分割和語義分割，及相關演算法思想。

使用EasyCV預測分割圖

EasyCV提供了在coco資料集上訓練的例項分割模型和全景分割模型以及在ADE20K上訓練的語義分割模型，參考EasyCV quick start（）完成依賴環境的配置後，可以直接使用這些模型完成對影像的分割譜預測，相關模型連結在reference中給出。

例項分割預測

由於該示例中的mask2fromer演算法使用了Deformable attention (在DETR系列演算法中使用該運算元可以有效提升演算法收斂速度和計算效率)，需要額外對該運算元進行編譯

cd thirdparty/deformable_attention
python setup.py build install

透過Mask2formerPredictor預測影像例項分割圖

import cv2
from easycv.predictors.segmentation import Mask2formerPredictor
predictor = Mask2formerPredictor(model_path='mask2former_instance_export.pth',task_mode='instance')
img = cv2.imread('000000123213.jpg')
predict_out = predictor(['000000123213.jpg'])
instance_img = predictor.show_instance(img, **predict_out[0])
cv2.imwrite('instance_out.jpg',instance_img)

輸出結果如下圖：

全景分割預測

透過Mask2formerPredictor預測影像全景分割圖

import cv2
from easycv.predictors.segmentation import Mask2formerPredictor
predictor = Mask2formerPredictor(model_path='mask2former_pan_export.pth',task_mode='panoptic')
img = cv2.imread('000000123213.jpg')
predict_out = predictor(['000000123213.jpg'])
pan_img = predictor.show_panoptic(img, **predict_out[0])
cv2.imwrite('pan_out.jpg',pan_img)

輸出結果如下圖：

語義分割預測

透過Mask2formerPredictor預測影像語義分割圖

import cv2
from easycv.predictors.segmentation import Mask2formerPredictor
predictor = Mask2formerPredictor(model_path='mask2former_semantic_export.pth',task_mode='semantic')
img = cv2.imread('000000123213.jpg')
predict_out = predictor(['000000123213.jpg'])
semantic_img = predictor.show_panoptic(img, **predict_out[0])
cv2.imwrite('semantic_out.jpg',semantic_img)

示例圖片來源：

在阿里雲機器學習平臺PAI上使用Mask2Former模型

PAI-DSW（Data Science Workshop）是阿里雲機器學習平臺PAI開發的雲上IDE，面向各類開發者，提供了互動式的程式設計環境。在DSW Gallery中( 連結 )，提供了各種Notebook示例，方便使用者輕鬆上手DSW，搭建各種機器學習應用。我們也在DSW Gallery中上架了Mask2Former進行影像分割的Sample Notebook（見下圖），歡迎大家體驗！

Mask2Former演算法解讀

上述例子中採用的模型是基於Mask2former實現的，Mask2former是一個統一的分割架構，能夠同時進行語義分割、例項分割以及全景分割，並且取得SO他的結果，在COCO資料集上全景分割精度57.8 PQ，例項分割精度達50.1 AP，在ADE20K資料集上語義分割精度達57.7 mIoU。

核心思想

Mask2Former採用mask classification的形式來進行分割，即透過模型去預測一組二值mask再組合成最終的分割圖。每個二值mask可以代表類別或例項，就可以實現語義分割、例項分割等不同的分割任務。

在mask classsification任務中，一個比較核心的問題是如何去找到一個好的形式學習二值Mask。如先前的工作 Mask R-CNN透過bounding boxes來限制特徵區域，在區域內預測各自的分割譜。這種方式也導致Mask R-CNN只能進行例項分割。Mask2Former參考DETR的方式，透過一組固定數量的特徵向量(object query)去表示二值Mask，透過Transformer Decoder進行解碼去預測這一組Mask。（ps：關於DETR的解讀可以參考：）

在DETR系列的演算法中，有一個比較重要的缺陷是在Transformer Decoder中的cross attention中會對全域性的特徵進行處理，導致模型很難關注到真正想要關注的區域，會降低模型的收斂速度和最終的演算法精度。對於這個問題Mask2former提出了Transformer Decoder with mask attention，每個Transformer Decoder block 會去預測一個attention mask並以0.5為閾值進行二值化，然後將這個attentino mask作為下一個block的輸入，讓attention模組計算時只關注在mask的前景部分。

模型結構

Mask2Former由三個部分組成：

Backbone（ResNet、Swin Transformer）從圖片中抽取低解析度特徵
Pixel Decoder 從低分辯率特徵中逐步進行上取樣解碼，獲得從低解析度到高解析度的特徵金字塔，迴圈的作為Transformer Decoder中V、K的輸入。透過多尺度的特徵來保證模型對不同尺度的目標的預測精度。

其中一層的Trasformer程式碼如下所示（ps：為了進一步加速模型的收斂速度，在Pixel Decoder中採用了Deformable attention模組）：

class MSDeformAttnTransformerEncoderLayer(nn.Module):
    def __init__(self,
                 d_model=256,
                 d_ffn=1024,
                 dropout=0.1,
                 activation='relu',
                 n_levels=4,
                 n_heads=8,
                 n_points=4):
                     super().__init__()
                     # self attention
                     self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
                     self.dropout1 = nn.Dropout(dropout)
                     self.norm1 = nn.LayerNorm(d_model)
                     # ffn
                     self.linear1 = nn.Linear(d_model, d_ffn)
                     self.activation = _get_activation_fn(activation)
                     self.dropout2 = nn.Dropout(dropout)
                     self.linear2 = nn.Linear(d_ffn, d_model)
                     self.dropout3 = nn.Dropout(dropout)
                     self.norm2 = nn.LayerNorm(d_model)
    @staticmethod
    def with_pos_embed(tensor, pos):
        return tensor if pos is None else tensor + pos
    def forward_ffn(self, src):
        src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))
        src = src + self.dropout3(src2)
        src = self.norm2(src)
        return src
    def forward(self,
                src,
                pos,
                reference_points,
                spatial_shapes,
                level_start_index,
                padding_mask=None):
                    # self attention
                    src2 = self.self_attn(
                        self.with_pos_embed(src, pos), reference_points, src,
                        spatial_shapes, level_start_index, padding_mask)
                    src = src + self.dropout1(src2)
                    src = self.norm1(src)
                    # ffn
                    src = self.forward_ffn(src)
                    return src

Transformer Decoder with mask attention 透過Object query和Pixel Decoder中得到的Multi-scale feature去逐層去refine二值mask圖，得到最終的結果。

其中核心的mask cross attention，會將前一層的預測的mask作為MultiheadAttention的atten_mask輸入，以此來將注意力的計算限制在這個query關注的前景中。具體實現程式碼如下：

class CrossAttentionLayer(nn.Module):
    def __init__(self,
                 d_model,
                 nhead,
                 dropout=0.0,
                 activation='relu',
                 normalize_before=False):
        super().__init__()
        self.multihead_attn = nn.MultiheadAttention(
            d_model, nhead, dropout=dropout)
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before
        self._reset_parameters()
    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos
    def forward_post(self,
                     tgt,
                     memory,
                     memory_mask: Optional[Tensor] = None,
                     memory_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None,
                     query_pos: Optional[Tensor] = None):
        tgt2 = self.multihead_attn(
            query=self.with_pos_embed(tgt, query_pos),
            key=self.with_pos_embed(memory, pos),
            value=memory,
            attn_mask=memory_mask,
            key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout(tgt2)
        tgt = self.norm(tgt)
        return tgt
    def forward_pre(self,
                    tgt,
                    memory,
                    memory_mask: Optional[Tensor] = None,
                    memory_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None,
                    query_pos: Optional[Tensor] = None):
        tgt2 = self.norm(tgt)
        tgt2 = self.multihead_attn(
            query=self.with_pos_embed(tgt2, query_pos),
            key=self.with_pos_embed(memory, pos),
            value=memory,
            attn_mask=memory_mask,
            key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout(tgt2)
        return tgt
    def forward(self,
                tgt,
                memory,
                memory_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, memory_mask,
                                    memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, memory_mask,
                                 memory_key_padding_mask, pos, query_pos)

Tricks

1.efficient multi-scale strategy

在pixel decoder中會解碼得到尺度為原圖1/32、1/16、1/8的特徵金字塔依次作為對應transformer decoder block的K、V的輸入。參照deformable detr的做法，對每個輸入都加上了sinusoidal positional embedding和learnable scale-level embedding。按解析度從低到高的循序依次輸入，並迴圈L次。

2.PointRend

透過PointRend的方式來節省訓練過程中的記憶體消耗，主要體現在兩個部分a.在使用匈牙利演算法匹配預測mask和真值標籤時，透過均勻取樣的K個點集代替完整的mask圖來計算match cost b.在計算損失時按照importance sampling策略取樣的K個點集代替完整的mask圖來計算loss（ps實驗證明基於pointreind方式來計算損失能夠有效提升模型精度）

3.Optimization improvements

更換了self-attention和cross-attention的順序。self-attention->cross-attention變成cross-attention->self-attention。
讓query變成可學習的引數。讓query進行監督學習可以起到類似region proposal的作用。透過實驗可以證明可學習的query可以產生mask proposal。
去掉了transformer deocder中的dropout操作。透過實驗發現這個操作會降低精度。

復現精度

例項分割及全景分割在COCO上的復現精度，實驗在單機8卡A100環境下進行(ps :關於例項分割復現精度問題在官方repo issue 46中有提及)

語義分割在ADE20K資料集上進行復現

使用EasyCV訓練分割模型

對於特定場景的分割，可以使用EasyCV框架和相應資料訓練定製化的分割模型。這裡以例項分割為例子，介紹訓練流程。

一、資料準備

目前EasyCV支援COCO形式的資料格式，我們提供了示例COCO資料用於快速走通流程。

wget 
mkdir -p data/  && mv small_coco_demo data/coco

二、模型訓練

在EasyCV的config資料夾下，我們提供了mask2former的資料處理和模型訓練及驗證的配置檔案(configs/segmentation/mask2former/mask2former_r50_8xb2_e50_instance.py),根據需要修改預測的類別、資料路徑。

執行訓練命令，如下所示：

#單機八卡
python -m torch.distributed.launch --nproc_per_node=8 --master_port 11111 tools/train.py \
                                        configs/segmentation/mask2former/mask2former_r50_8xb2_e50_instance.py \
                                        --launcher pytorch \
                                        --work_dir experiments/mask2former_instance \
                                        --fp16

模型匯出，將config檔案儲存到模型中，以便在predictor中得到模型和資料處理的配置，匯出後的模型就可直接用於分割圖的預測。

python tools/export.py configs/segmentation/mask2former/mask2former_r50_8xb2_e50_instance.py epoch_50.pth mask2former_instance_export.pth

END

EasyCV會持續進行SOTA論文復現進行系列的工作介紹，歡迎大家關注和使用，歡迎大家各種維度的反饋和改進建議以及技術討論，同時我們十分歡迎和期待對開源社群建設感興趣的同行一起參與共建。

使用EasyCV Mask2Former輕鬆實現影像分割

導言