SegFix原理及源码--2020.8

论文SegFix: Model-Agnostic Boundary Refinement for Segmentation

本文提出了一种模型无关的后处理方案，即用内部像素的预测代替原来不可靠的边界像素预测，以提高由任何现有分割模型生成的分割结果的边界质量。该方法仅对输入图像进行两步处理：(i)定位边界像素；(ii)识别每个边界像素对应的内部像素（通过学习从边界像素到内部像素的方向来建立对应关系）。

该方法不需要先验信息的分割模型，达到接近实时的速度。实验验证，SegFix减少了cityspace数据集中各种先进模型(ADE20K和GTA5)所生成的分割结果的边界误差。

1、原理

1.1 SegFix原理

提出了一种与模型无关的后处理机制来降低边界的分类错误。

方法是边界的类别用相对应的内部pixel的类别来表示。

而相对应的内部pixel可以通过以下方法来确定：

1）首先要先确定物体边界：用卷积来预测一个binary mask来表示边界。1表示边界，0表示other
2）学习一个从边界pixel到内部pixel的方向
3）让边界pixel沿着这个方向移动特定的距离即可确定该内部pixel。

可以用下面这个式子来说明：

$p_i$表示边界的pixel，$\Delta q_i$表示指向内部pixel的偏移向量offset vector。$\tilde{L}$表示refined label map，
如果边界很厚，有可能会有一些假的内部pixel，我们用一个因子来rescale这个偏移向量。
只需要产生一次offset map，就可以将该offset map应用到任何分割模型的预测结果中做refine。

1.2 SegFix结构

• Boundry branch

$1\times 1\times 256$ Conv-BN-ReLU + $1\times 1\times 2$ Conv+Upsample得到boundry_map
boundry_map是一个binary mask，1表示边界，0表示其它。
loss是binary cross-entropy loss, 权重为1。

• Direction branch

$1\times 1\times 256$ Conv-BN-ReLU + $1\times 1\times m$ Conv+Upsample得到direction_map
预测所有位置的像素与之最近的同类像素的方向，这里的方向不是连续的，而是对[0,2π)量化成m（m=8）个值后的结果；
将Direction map和Boundry map做element-wise product，即只在边界处计算loss。
loss是category cross-entropy loss。权重为1

• offset branch

用来转换预测得到的D（已经乘过了B ）到坐标偏移图offset_map，$\Delta Q\in R^{H\times W\times 2}$
不同的方向会被映射到不同的坐标偏移值。例如下图中，对于m=4的右上方向，可以映射为(1, 1) 最终通过坐标映射（the grid-sample scheme），重新对现有方法在边界区域的预测结果进行调整。

使用$\Delta Q\in R^{H\times W\times 2}$调整现有粗糙预测$L$的公式为：${\tilde{L}}_{p_i}=L_{p_i+\Delta q_i}$，即将每个位置上最终的预测，设定为坐标偏移后对应位置的原始预测结果

上图中，不同的方向会被映射到不同的坐标偏移值。如m=4个方向范畴对应于偏移量(1,1)、(−1,1)、(−1，−1)和(1，−1)。对于m=4的右上方向映射为(1, 1)；

1.3 训练Training

在训练阶段，首先将图像输入到主干网络中来预测特征图，然后应用边界分支预测二值边界图和应用方向分支预测方向图，我们将边界损失和方向损失分别应用到预测的边界图和方向图上。

1.4 测试Testing生成offset

在测试阶段，我们首先将模型执行到图像上，以生成偏移量映射，然后根据偏移量图对现有方法的分割结果进行细化。

• 生成offset

def _get_offset(self, mask_logits, dir_logits):

        edge_mask = mask_logits[:, 1] > 0.5
        dir_logits = torch.softmax(dir_logits, dim=1)
        n, _, h, w = dir_logits.shape

        keep_mask = edge_mask

        dir_label = torch.argmax(dir_logits, dim=1).float()
        offset = DTOffsetHelper.label_to_vector(dir_label)
        offset = offset.permute(0, 2, 3, 1)
        offset[~keep_mask, :] = 0
        return offset

DTOffsetHelper.label_to_vector(dir_label)

def label_to_vector(labelmap, 
                        num_classes=DTOffsetConfig.num_classes):

        assert isinstance(labelmap, torch.Tensor)

        mapping = label_to_vector_mapping[num_classes]
        offset_h = torch.zeros_like(labelmap).long()
        offset_w = torch.zeros_like(labelmap).long()

        for idx, (hdir, wdir) in enumerate(mapping):
            mask = labelmap == idx
            offset_h[mask] = hdir
            offset_w[mask] = wdir

        return torch.stack([offset_h, offset_w], dim=-1).permute(0, 3, 1, 2).to(labelmap.device)

label_to_vector_mapping[num_classes]

label_to_vector_mapping = {
    4: [
        [-1, -1], [-1, 1], [1, 1], [1, -1]
    ] if not DTOffsetConfig.c4_align_axis else [
        [0, -1], [-1, 0], [0, 1], [1, 0]
    ],    
    8: [
        [0, -1], [-1, -1], [-1, 0], [-1, 1],
        [0, 1], [1, 1], [1, 0], [1, -1]
    ],
    16: [
        [0, -2], [-1, -2], [-2, -2], [-2, -1], 
        [-2, 0], [-2, 1], [-2, 2], [-1, 2],
        [0, 2], [1, 2], [2, 2], [2, 1],
        [2, 0], [2, -1], [2, -2], [1, -2]
    ],
    32: [
        [0, -4], [-1, -4], [-2, -4], [-3, -4], [-4, -4], [-4, -3], [-4, -2], [-4, -1],
        [-4, 0], [-4, 1], [-4, 2], [-4, 3], [-4, 4], [-3, 4], [-2, 4], [-1, 4],
        [0, 4], [1, 4], [2, 4], [3, 4], [4, 4], [4, 3], [4, 2], [4, 1],
        [4, 0], [4, -1], [4, -2], [4, -3], [4, -4], [3, -4], [2, -4], [1, -4],
    ]
}

考虑到边界较厚时可能你会有一些“伪”内部像素，提出了两种不同的方案：

（1）将所有偏移量重新缩放一个因子；

（2）迭代的应用偏移量（“伪造”内部像素的）直到找到内部像素。

1.5 Refinement对结果进行细化

scripts.cityscapes.segfix.py

offset_map = get_offset(basename) # 读取保存的offset的mat文件
output_label_map = shift(input_label_map, offset_map)

def get_offset(basename):
    return io.loadmat(osp.join(offset_dir, basename+'.mat'))['mat']\
        .astype(np.float32).transpose(2, 0, 1) * args.scale

def shift(x, offset):
    """
    x: h x w
    offset: 2 x h x w
    """
    h, w = x.shape
    x = torch.from_numpy(x).unsqueeze(0)
    offset = torch.from_numpy(offset).unsqueeze(0)
    coord_map = gen_coord_map(h, w)
    norm_factor = torch.FloatTensor([(w-1)/2, (h-1)/2])
    grid_h = offset[:, 0]+coord_map[0]
    grid_w = offset[:, 1]+coord_map[1]
    grid = torch.stack([grid_w, grid_h], dim=-1) / norm_factor - 1
    x = F.grid_sample(x.unsqueeze(1).float(), grid, padding_mode='border', mode='bilinear', align_corners=True).squeeze().numpy()
    x = np.round(x)
    return x.astype(np.uint8)

最终通过坐标映射（the grid-sample scheme ），重新对现有方法在边界区域的预测结果进行调整，过程如下图所示。

基于偏移映射对粗略标记图进行细化。用不同的箭头表示不同的偏移向量。在粗略标签图中红色标记错误位置，在精确标签图中箭头方向标记相应的修正位置。

2、ground-truth

先从ground-truth分割图生成distance map，再在此基础上生成boundary map和direction map。

2.1 Distance map

利用传统的距离转换（distance transform）方法，需要先生成一个distance map，表示当前像素到其它类别中像素的最小距离（这里用的是欧式距离）。

在文中分了三步来做，最终得到的是每个像素到其边界的距离。

1）先将ground-truth分成K个binary map，即每个类别一个map，1表示当前类别，0表示其它；
2）在每个binary_map上单独计算相距属于其他类别像素的最小欧式距离，得到每个类别的distance map，其实就是每个像素到其边界的距离。 这里使用了scipy的函数scipy.ndimage.morphology.distance_transform_edt()。这个函数用于距离转换，计算图像中非零点到最近背景点（即0）的距离。它被用在每个类别独立的binary_mask上，正好计算的就是相距于其他类别（每个binary mask上的0表示的就是“其他类别”）的最小欧氏距离
3）再将这K个distance_map混合，得到最终的fuse distance map。

segfix.lib.datasets.preprocess.cityscapes.dt_offset_generator.py

def process(inp):
    (indir, outdir, basename) = inp
    print(inp)
    labelmap = np.array(Image.open(osp.join(indir, basename)).convert("P")).astype(np.int16)
    labelmap = _encode_label(labelmap)
    labelmap = labelmap + 1 # 因为要用0表示非当前类别，所以这里标签从1开始。
    depth_map = np.zeros(labelmap.shape, dtype=np.float32)
    dir_map = np.zeros((*labelmap.shape, 2), dtype=np.float32)

    for id in range(1, len(label_list) + 1):
        labelmap_i = labelmap.copy()
        labelmap_i[labelmap_i != id] = 0
        labelmap_i[labelmap_i == id] = 1 # 1.每个类别一个label_map

        if labelmap_i.sum() < 100:
            continue

        if args.metric == 'euc': # 用欧式距离
            depth_i = scipy.ndimage.morphology.distance_transform_edt(labelmap_i) # 2.每个像素到其边界的距离,distance_map_i
        elif args.metric == 'taxicab':
            depth_i = distance_transform_cdt(labelmap_i, metric='taxicab')
        else:
            raise RuntimeError
        depth_map += depth_i # 3.多个类别的混合，得到最终的distance_map
        
        dir_i_before = dir_i = np.zeros_like(dir_map)
        # 4.direction_map用sobel算子
        dir_i = torch.nn.functional.conv2d(torch.from_numpy(depth_i).float().view(1, 1, *depth_i.shape), sobel_ker, padding=ksize//2).squeeze().permute(1, 2, 0).numpy() 

        # The following line is necessary
        dir_i[(labelmap_i == 0), :] = 0 
        
        dir_map += dir_i # 5.获取direction_map

    depth_map[depth_map > 250] = 250 # 这里为了后面的阈值比较，距离小于5的是边界像素
    depth_map = depth_map.astype(np.uint8)
    deg_reduce = 2
    dir_deg_map = np.degrees(np.arctan2(dir_map[:, :, 0], dir_map[:, :, 1])) + 180 # 0~360
    dir_deg_map = (dir_deg_map / deg_reduce) # 0~180
    print(dir_deg_map.min(), dir_deg_map.max())
    dir_deg_map = dir_deg_map.astype(np.uint8) 

    io.savemat(
        osp.join(outdir, basename.replace("png", "mat")),
        {"dir_deg": dir_deg_map, "depth": depth_map, 'deg_reduce': deg_reduce},
        do_compression=True,
    )
    
    try:
        io.loadmat(osp.join(outdir, basename.replace("png", "mat")),)
    except Exception as e:
        print(e)
        io.savemat(
            osp.join(outdir, basename.replace("png", "mat")),
            {"dir_deg": dir_deg_map, "depth": depth_map, 'deg_reduce': deg_reduce},
            do_compression=False,
        )

2.2 Boundary map

将上面得到的fuse distance map中，距离小于某个阈值$\gamma$($\gamma =5$)的作为边界区域的像素，标记为1，大于阈值的认为在特定目标区域内部，标记为0。因为距离值越小，说明越接近边界.

源码如下，其中

DTOffsetConfig.max_distance=5
DTOffsetConfig.min_distance=0

keep_mask = (distance_map <= DTOffsetConfig.max_distance) & (distance_map >= DTOffsetConfig.min_distance)
mask_label_map[keep_mask] = 1
mask_label_map[seg_label_map == -1] = -1

2.3 Direction map

direction map的计算是在未合并的K个单独的distance map上，分别使用$9\times 9$的Sobel滤波器。基于Sobel滤波器的方向是在[0°，360°)内，并且每个像素位置的方向都指向邻域内部距离目标边界最远的像素。

整个方向范围被均匀划分成m=8类，然后每个像素的方向被赋值成对应的方向类别。

def align_angle(angle_map, 
                    num_classes=DTOffsetConfig.num_classes, 
                    return_tensor=False):

        if num_classes == 4 and not DTOffsetConfig.c4_align_axis:
            return DTOffsetHelper.align_angle_c4(angle_map, return_tensor=return_tensor)

        if return_tensor:
            assert isinstance(angle_map, torch.Tensor)
        else:
            assert isinstance(angle_map, np.ndarray)

        step = 360 / num_classes
        if return_tensor:
            new_angle_map = torch.zeros(angle_map.shape).float().to(angle_map.device)
            angle_index_map = torch.zeros(angle_map.shape).long().to(angle_map.device)
        else:
            new_angle_map = np.zeros(angle_map.shape, dtype=np.float)
            angle_index_map = np.zeros(angle_map.shape, dtype=np.int)
        mask = (angle_map <= (-180 + step/2)) | (angle_map > (180 - step/2))
        new_angle_map[mask] = -180
        angle_index_map[mask] = 0

        for i in range(1, num_classes):
            middle = -180 + step * i
            mask = (angle_map > (middle - step / 2)) & (angle_map <= (middle + step / 2))
            new_angle_map[mask] = middle
            angle_index_map[mask] = i

        return new_angle_map, angle_index_map

上述代码的计算结果就如下图所示

整体过程如下图所示：

上图表示二值映射→距离映射→方向映射。ground-truth二值映射分为car、road和side-walk三类。首先对每个二值映射应用距离变换来计算真值距离映射。然后在距离图上使用Sobel滤波器计算真值方向图。选择不同的颜色来表示不同的距离值或方向值。