SOLO算法笔记

本文为个人学习过程中所记录笔记，便于梳理思路和后续查看用，如有错误，感谢批评指正！

paper:SOLO Segmenting Objects by Locations
code:https://github.com/open-mmlab/mmdetection/blob/main/mmdet/models/dense_heads/solo_head.py

abstract

  通常做实例分割：（1）先检测，再分割；（2）首先预测嵌入向量，然后使用聚类技术将像素分组到单独的实例中。
  本文：通过引入实例类别的概念，即根据实例的位置和大小为实例中的每个像素分配类别，从而很好地将实例分割转换为单一的分类可解问题。

Introduction

  图像中的目标属于一组固定的语义类别，但实例的数量是不同的。因此，语义分割可以很容易地表示为密集的单像素分类问题，而直接按照相同的范例预测实例标签是具有挑战性的。
  当前实例分割可分为两类范式：自上而下（先检测后分割，过于依赖检测效果）和自下而上范式（通过学习一个同性关系，为每个像素分配一个嵌入向量，聚合同一个实例像素，疏远不同实例像素，依赖像素级嵌入向量的学习和分组处理方法）。
  作者重新考虑实例分割问题，分析coco数据集，作者认为可以通过不同实例间的距离和各个实例的尺寸大小来分割区分各个实例。
  目前较为流行的范式是采用全卷积网络，输出N张特征图，其中每维度的特征图代表一个类别，包括背景。类似地，本文将引入一个实例类别的概念，例如量化的中心位置和目标尺寸将用于分割目标。
位置：一个图分成S x S个格点，代表S * S个位置类别。每一个输出通道一个中心位置类别。这里没有太看懂，需要结合代码。

  因此，通过将每个像素分类为其实例位置类别，位置预测问题转换成分类问题，更直接的对于数量变化的实例个数进行建模，而不需要依赖类似于分组和学习嵌入向量这种后处理方法。
尺寸：采用FPN去分辨不同目标尺寸。

Related Work

自上而下：
  （1）与mask rcnn相比，PANET加强了特征表示
  （2）Mask Scoring R-CNN增加了mask-iou分支去量化mask质量。
自下而上：
  （1）SGN
  （2）SSAP
直接实例分割：
  据作者所知，暂时没有，一些方法被看作半直接范式：AdaptIS，POlarMask，

2 Our Method: SOLO

2.1 Problem Formulation
  solo的主要思想是将实例分割看作是类别感知预测问题。具体地，将输入图片划分成sxs个格子，如果某个目标的中心落入某个格子，这个格子将负责该目标的语义类别预测和目标实例的分割两个任务。
  对于每个grid都有c维输出，表示c个类别的概率。输出为S x S x C如图所示：将实例分割看成是两个子任务，类别预测和实例mask的生成。该假设基于每个grid对应一个单独的实例，也就是每个grid里仅有一个物体中心点。

Instance Mask：对这些掩码在第三维度进行显式编码，得到一个3D输出向量。mask输出维度为H_I x W_I x S²。第k个通道负责分割第（i，j）个grid。k= i * S + j。由此将类别预测和mask预测联系起来。
  可以采用全卷积网路FCNs直接对实例掩码进行提取，但是FCNs具有空间不变性，这对于分类、语义分割等任务很好，但是我们希望模型对位置具有敏感性，因为我们的分割掩模是以网格单元为条件的，并且必须由不同的特征通道来分隔。
  我们受启于CoordConv，在网络的开始变加入归一化的像素坐标。具体地，创建一个和输入大小相同的tensor，以归一化到[-1,1]之间的像素坐标填充。特征尺寸将从H X W X D变为H X W X (D+2)，最后两个通道是（x, y）像素坐标。代码如下：

def generate_coordinate(featmap_sizes, device='cuda'):
    """Generate the coordinate.

    Args:
        featmap_sizes (tuple): The feature to be calculated,
            of shape (N, C, W, H).
        device (str): The device where the feature will be put on.
    Returns:
        coord_feat (Tensor): The coordinate feature, of shape (N, 2, W, H).
    """

    x_range = torch.linspace(-1, 1, featmap_sizes[-1], device=device) #-1到1之间
    y_range = torch.linspace(-1, 1, featmap_sizes[-2], device=device)
    y, x = torch.meshgrid(y_range, x_range)
    y = y.expand([featmap_sizes[0], 1, -1, -1])
    x = x.expand([featmap_sizes[0], 1, -1, -1])
    coord_feat = torch.cat([x, y], 1)

    return coord_feat

实例分割建模：最后需要采用nms。
2.2 Network Architecture
  SOLO采用卷积backbone进行特征提取，采用FPN生成特征金字塔融合不同尺寸。这些特征图作为语义类别和实例mask的输入，不同尺度的输入共享head权重，不同尺度的grid数目不同，因此最后一层卷积权重不共享。head结构如图：

2.3 SOLO Learning
Label Assignment：类别预测分支，如果有真值中心落入，就是正样本，否则为负样本。采用较火的中心采样方法。控制尺度，让每个真值具有平均3个正样本。

  二进制分割mask对应每个正样本。每张图S²个grid，所以有S²个输出mask。

Loss Function:
  损失函数包括两部分：
Locate表示用于语义分类的focal loss，Lmask用于mask预测。
d_mask采用BCEloss，focal loss和dice loss都做了对比。最后采用了DICE loss，拉梅达设置为3。

2.4 Inference

  推理过程中，直接前向推理得出各个格点的概率分数以及mask。首先设置一个0.1的阈值筛除。然后选取前500个分数的mask进行nms操作。采用0.5的阈值转换mask成二进制mask。
  每个预测的分类分数乘以mask作为最终置信度分数。
  具体结果如下：

def forward(self, x: Tuple[Tensor]) -> tuple:
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of classification scores and mask prediction.

                - mlvl_mask_preds (list[Tensor]): Multi-level mask prediction.
                  Each element in the list has shape
                  (batch_size, num_grids**2 ,h ,w).
                - mlvl_cls_preds (list[Tensor]): Multi-level scores.
                  Each element in the list has shape
                  (batch_size, num_classes, num_grids ,num_grids).
        """
        assert len(x) == self.num_levels #输入为多个FPN多个尺度的特征图
        feats = self.resize_feats(x)     #将第一个特征图下采样两倍，最后一个特征图上采样然后append成list
        mlvl_mask_preds = []
        mlvl_cls_preds = []
        for i in range(self.num_levels):
            x = feats[i]
            mask_feat = x
            cls_feat = x
            # generate and concat the coordinate
            coord_feat = generate_coordinate(mask_feat.size(),
                                             mask_feat.device) #(N, 2, W, H)
            mask_feat = torch.cat([mask_feat, coord_feat], 1) #(N, D + 2, H, W)

            for mask_layer in (self.mask_convs):
                mask_feat = mask_layer(mask_feat)

            mask_feat = F.interpolate(
                mask_feat, scale_factor=2, mode='bilinear')
            mask_preds = self.conv_mask_list[i](mask_feat) #一通卷积操作下来，最后输出通道数为S^2

            # cls branch
            for j, cls_layer in enumerate(self.cls_convs):
                if j == self.cls_down_index:
                    num_grid = self.num_grids[i]
                    cls_feat = F.interpolate(
                        cls_feat, size=num_grid, mode='bilinear')
                cls_feat = cls_layer(cls_feat)

            cls_pred = self.conv_cls(cls_feat) #多个卷积操作下来以后，最后通过卷积输出num_class个通道

            if not self.training:
                feat_wh = feats[0].size()[-2:] #第一个特征图的wh
                upsampled_size = (feat_wh[0] * 2, feat_wh[1] * 2) #扩充两倍
                mask_preds = F.interpolate(
                    mask_preds.sigmoid(), size=upsampled_size, mode='bilinear') #上采样
                cls_pred = cls_pred.sigmoid() #激活函数
                # get local maximum
                local_max = F.max_pool2d(cls_pred, 2, stride=1, padding=1) #maxpooling
                keep_mask = local_max[:, :, :-1, :-1] == cls_pred
                cls_pred = cls_pred * keep_mask

            mlvl_mask_preds.append(mask_preds) #(batch_size, num_grids**2 ,h ,w)
            mlvl_cls_preds.append(cls_pred) #(batch_size, num_classes, num_grids ,num_grids).
        return mlvl_mask_preds, mlvl_cls_preds

3.2 How SOLO Works?
  设置S=12，
  grid的影响：

  多尺度融合提升巨大，图2 可以看出
  CoordConv采用带来提升：

  损失函数影响：

  误差分析证明mask分支还有极大改进的空间

4 Decoupled SOLO
  mask分支采用Decoupled SOLO

`