使用 EasyCV Mask2Former 轻松实现图像分割

导言

图像分割 (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

1. 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 中有提及)

语义分割在 ADE20K 数据集上进行复现

使用 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