图像ROI区域如何映射到feature map—Mask R-CNN

问题描述

某些情况下需要将卷积/池化层后的features map的某些patch映射回原图像，就需要features map怎么变过来的怎么回去。

#read image
self.object_detector = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained = True, trainable_backbone_layers = 0)
images_nested = torchvision.models.detection.image_list.ImageList(self.object_detector.transform.normalize(images), targets['image_height_width_resized'].tolist())

这里试验的模型是two-stage 的Mask R-CNN，通过在每个感兴趣区域 (RoI) 上添加一个预测分割掩码的分支来扩展 Faster R-CNN，与现有的分类和边界框回归分支并行运算。

在 Faster RCNN 中，proposal generation 是 RPN(Region Proposal Network)，会根据 feature map 生成数十万个候选框，通过 NMS 选出前景概率最高的 2000 个框。由于候选框（特征图）的大小各异，通过 ROI pooling，得到固定大小的输出，channel 数量就是框的数量。ROI pooling 的特点是输入特征图尺寸不固定，但是输出特征图尺寸固定。最后经过全连接层得到回归和分类的输出。

box features

#self.backbone(images.tensors)返回的features是一个 dict，每个元素是一个 feature map，每个特征图的宽高是上一个特征图宽高的一半。
#这 n 个 feature map 分别对应 ResNet 中的 n个layer的特征图
img_features = self.object_detector.backbone(images_nested.tensors)
#通过ROI pooling branch，把不同尺度的特征图池化到相同尺度
box_features = self.object_detector.roi_heads.box_roi_pool(img_features, bbox.unbind(), images_nested.image_sizes)#

上面的代码中RPN生成的候选框由训练集中的boxes代替，省略RPN，经过box_roi_pooling后生成相同尺寸的特征图。

mask features

分割掩码的分支的features map变换类似：

#经过mask_roi_pooling后生成相同尺寸的特征图。
mask_features = self.object_detector.roi_heads.mask_roi_pool(img_features, bbox.unbind(), images_nested.image_sizes)
mask_features = self.object_detector.roi_heads.mask_head(mask_features)
#mask_logits输出最终的分类
mask_logits = self.object_detector.roi_heads.mask_predictor(mask_features)
#在已知category的条件下训练，直接索引mask的预测结果
mask_probs = self.index_select_batched(mask_logits, category_idx.flatten()).sigmoid()#索引batch

不同之处在于多了mask_head这个结构，在经过试验之后发现features map经过mask_head后shape并没有变化，下面贴出的是pytorch中的mask_head源码：

class MaskRCNNHeads(nn.Sequential):
    def __init__(self, in_channels, layers, dilation):
        """
        Args:
            in_channels (int): number of input channels
            layers (list): feature dimensions of each FCN layer
            dilation (int): dilation rate of kernel
        """
        d = OrderedDict()
        next_feature = in_channels
        for layer_idx, layer_features in enumerate(layers, 1):
            d["mask_fcn{}".format(layer_idx)] = nn.Conv2d(
                next_feature, layer_features, kernel_size=3,
                stride=1, padding=dilation, dilation=dilation)
            d["relu{}".format(layer_idx)] = nn.ReLU(inplace=True)
            next_feature = layer_features

        super(MaskRCNNHeads, self).__init__(d)
        for name, param in self.named_parameters():
            if "weight" in name:
                nn.init.kaiming_normal_(param, mode="fan_out", nonlinearity="relu")
            # elif "bias" in name:
            #     nn.init.constant_(param, 0)

mask_head的作用就是再次确定，映射后得到的bbox是否是真的有物体，并将它们的坐标再次进行细调。

映射关系

通过上面的流程梳理，会发现映射生在mask/box_roi_pool这一步，通过MultiScaleRoIAlign(这里使用的是RoiAlign)将不同尺寸的图像Roi区域变换成统一尺寸，所以只需要了解RoiAlign的原理就可以了。
例如：

最后得到mask_probs的shape为[1, 28, 28]，原始图片尺寸是 ： [395,244]，图片中Roi区域为 “box”: [ 4,22,362,228]
随机选取Roi区域内的“image patch-box”坐标为：[172, 24, 291, 93]
那么对于features map中的点（xf，yf）其在image patch中的坐标（xp，yp）结合RoiAlign的原理，通过相似性原理可以粗略得到：

xf*(291-172)/28+172=xp
yf*(93-24)/28+24=yp

参考文章：[PyTorch 学习笔记] 8.2 目标检测简介
一文读懂 RoIPooling、RoIAlign 和 RoIWarp