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

簡介： EasyCV可以輕鬆預測圖像的分割譜以及訓練定製化的分割模型。本文主要介紹如何使用EasyCV實現實例分割、全景分割和語義分割，及相關算法思想。

作者：賀弘 謙言 臨在

導言

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

使用EasyCV預測分割圖

EasyCV提供了在coco數據集上訓練的實例分割模型和全景分割模型以及在ADE20K上訓練的語義分割模型，參考EasyCV quick start（https://github.com/alibaba/EasyCV/blob/master/docs/source/quick_start.md）完成依賴環境的配置後，可以直接使用這些模型完成對圖像的分割譜預測，相關模型鏈接在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)

示例圖片來源：cocodataset

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

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

Mask2Former算法解讀

上述例子中採用的模型是基於Mask2former實現的，Mask2former是一個統一的分割架構，能夠同時進行語義分割、實例分割以及全景分割，並且取得SOTA的結果，在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的解讀可以參考：基於EasyCV復現DETR和DAB-DETR，Object Query的正確打開方式）

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

模型結構

Mask2Former由三個部分組成：

1. Backbone（ResNet、Swin Transformer）從圖片中抽取低分辨率特徵
2. 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

3. 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

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

復現精度

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

Model	PQ	Box mAP	Mask mAP	memory	train_time
mask2former_r50_instance_official			43.7
mask2former_r50_8xb2_epoch50_instance		46.09	43.26	13G	3day2h
mask2former_r50_panoptic_official	51.9		41.7
mask2former_r50_8xb2_epoch50_panoptic	51.64	44.81	41.88	13G	3day4h

 

語義分割在ADE20K數據集上進行復現

Model	mIoU	train memory	train_time
mask2former_r50_semantic_official	47.2
mask2former_r50_8xb2_e127_samantic	47.03	5.6G	15h35m

 

使用EasyCV訓練分割模型

對於特定場景的分割，可以使用EasyCV框架和相應數據訓練定製化的分割模型。這裏以實例分割爲例子，介紹訓練流程。

一、數據準備

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

wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_coco_demo/small_coco_demo.tar.gz && tar -zxf small_coco_demo.tar.gz

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

Reference

實例分割模型：http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/mask2former_r50_instance/mask2former_instance_export.pth

全景分割模型：http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/mask2former_r50_panoptic/mask2former_pan_export.pth

語義分割模型：http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/segmentation/mask2former_r50_semantic/mask2former_semantic_export.pth

EasyCV往期分享

EasyCV開源地址：https://github.com/alibaba/EasyCV

EasyCV DataHub 提供多領域視覺數據集下載，助力模型生產 https://zhuanlan.zhihu.com/p/572593950

EasyCV帶你復現更好更快的自監督算法-FastConvMAE https://zhuanlan.zhihu.com/p/566988235

基於EasyCV復現DETR和DAB-DETR，Object Query的正確打開方式 https://zhuanlan.zhihu.com/p/543129581

基於EasyCV復現ViTDet：單層特徵超越FPN https://zhuanlan.zhihu.com/p/528733299

MAE自監督算法介紹和基於EasyCV的復現 https://zhuanlan.zhihu.com/p/515859470

EasyCV開源｜開箱即用的視覺自監督+Transformer算法庫 https://zhuanlan.zhihu.com/p/50521999

END

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

原文鏈接：https://click.aliyun.com/m/1000364647/

本文爲阿里雲原創內容，未經允許不得轉載。