论文解读 CascadePSP: Toward Class-Agnostic and Very High-Resolution Segmentation via Global and Local Re

论文地址：https://arxiv.org/pdf/2005.02551.pdf
发表时间：2020
代码地址：https://github.com/hkchengrex/CascadePSP

2020年最先进的语义分割方法几乎只在一个固定的分辨率范围内的图像上进行训练，这对于非常高分辨率的图像是不准确的，因为使用低分辨率分割的双边上采样不能充分捕捉沿物体边界的高分辨率细节。在本文中，我们提出了一种新的方法——CascadePSP网络，来解决高分辨率分割问题，而不使用任何高分辨率训练数据。它可以尽可能地细化和纠正局部边界。虽然CascadePSP是用低分辨率的分割数据进行训练的，但CascadePSP适用于任何分辨率，即使是对于大于4K的非常高分辨率的图像。我们对不同的数据集进行了定量和定性的研究，以表明CascadePSP可以使用我们的新的重新细化模块来揭示像素精确的分割边界，而不需要进行任何微调。因此，我们的方法可以看作是类不可知论的方法。

关键总结

• 1、可用于提升任意语义分割模型的精度，本质上是一种用深度学习实现的后处理方法
• 2、在训练时需要输入其他模型的初步语义分割结果（实际上是对GT做扰动，使模型在精度提升上无偏好）
• 3、在2.3节中的loss消融实验可以发现，经过第一次提炼（使用ce loss输出步幅为8）时，在deeplabv3+上iou已经提升了1.6[最终提升1.8]
• 4、在2.3节中的结构消融实验可以发现，经过第三次提炼（使用输出步幅为1的输出）时，在deeplabv3+上iou已经提升了1.3[最终提升1.7]
• 5、CascadePSP在结果提升时是做了级联（输出步幅8,4,1）对应着Global step可以提示语义分割的主体精度，对于高分影像最有用的为Local step，也就是对图像进行crop然后预测
• 6、对于大部分的语义分割可以增强其结果。其仅使用Global step做增强的部分，可以很便捷的改写为c++模型；使用Local step做增强的部分需要，在c++上实现图像切片才可实现。

基于本论文代码，实现多值语义分割结果的伪标签可以参考https://blog.csdn.net/a486259/article/details/129559350

1、背景

1.1 现有方法在高分辨率图像中存在局限性

随着相机和显示器分辨率的提升，高分影像的精准分割面临着技术挑战。现有的语义分割技术都是在PASCAL或者COCO上研究的，这些属于低分辨率影像；此外，低分影像模型推广到高分影像需要重新训练，而图像size的增大，使的gpu需求成平方倍扩展。故此，基于此研究出的模型很难不能适用到4k的高分辨率影像中。 此外，缺乏高分辨率的图像分割数据集是评价高分辨率图像分割模型的困难问题之一。

1.2 高分影像识别的难点

训练数据量

高分影像语义分割，需要像素级的标注，对数据标注成本较大，很难获得像素级的标注

语义感受野

在高分影像上训练模型需要更大的语义感受野，且增加了大size目标域小size目标的视野差距（在低分影像中大size目标显示不完全，故限制了目标size差异）。目前使用降采样或图像裁剪来降低input size，但是降采样删除了图像细节，图像裁剪破坏了图像上下文（可能将大目标截断）。

2、方法

本文提出了CascadePSP，一种通用的分割细化模型，可以将任何给定的分割从低分辨率细化到高分辨率。CascadePSP是独立训练的，可以附加到任何现有的方法来改进其分割，可以产生一个更精细、更准确的对象分割掩模。CascadePSP将一个初始掩码作为输入，它可以作为任何算法的输出，以提供一个粗略的对象位置。然后CascadePSP将输出一个改进的掩码。CascadePSP以级联方式设计，以粗到细的方式生成精细的分割。来自早期级别的粗输出预测对象结构，这将被用作后者级别的输入，以细化边界细节。下图展示了CascadePSP在非常高的分辨率下，也能细化和纠正错误的边界，以产生更准确的结果。

模型不需要针对特定的数据集或特定模型的输出进行针对训练，在不进行微调的情况下可以提高最先进的分割模型的性能；
可以通过通过扰动GT-seg值来模拟其他模型的pred-seg结果训练模型，来增强模型性能，方法可以直接扩展到密集场景语义分割；
在本文的研究中，提出了用于高分影像分割评估的数据集，BIG。

2.1 模型结构

如图2所示，我们的细化模块采用不同尺度上的图像和多个不完全分割掩模来生成细化的分割。多尺度输入允许模型捕获不同层次的结构和边界信息，这允许网络学习自适应地融合来自不同的掩模特征尺度在最好的水平上细化分割。

在图2中所示的结构中，所有具有较低分辨率的seg结果输入都被双向上采样到相同的大小，并与RGB图像相连接。使用以ResNet-50为骨干的PSPNet输入中提取步幅8个特征图。遵循PSPNet中设定的金字塔池大小[1,2,3,6]，这有助于捕获全局上下文。除了最终的步幅1输出外，CascadePSP还生成了中间步幅8和步幅4的seg结果[`跳过了输出步幅为2的输出，以提供纠正局部误差边界的灵活性`]，这可以使输入分割的整体结构变得更稳定。

为了重构在提取过程中丢失的像素级图像细节，CascadePSP采用了来自主干网络的跳跃连接，并使用一个上采样块来融合特征。CascadePSP将跳跃连接特征和主分支的双上采样特征连接起来，并用两个ResNet block对它们进行处理。分割输出使用2层1×1 conv生成，然后是sigmoid。

2.2 loss设计

对于不同的步幅使用不同的损失函数，因为粗细化关注全局结构而忽略局部细节，而精细细化则是通过依赖局部线索实现像素级的精度。

• 对于粗尺度，步幅8输出使用交叉熵损失，
• 对于细尺度，步幅1输出使用L1+L2损失，
• 对于中间尺度，步幅4输出使用交叉熵和L1+L2损失的平均值
  完整的loss函数设计如下
  

此外，为了鼓励更好的边界细化，对步幅1的输出采用了分割梯度幅度上的L1损失。分割梯度由3×3平均滤波器和Sobel算子估计。梯度损失使输出在像素级上更好地坚持对象边界。由于梯度与像素级损失相比更稀疏，我们用α进行权衡，在我们的实验中，α将其设置为5。梯度损失可以写成以下形式，其中$f_m()$表示3x3的均值滤波，∇表示Sobel算子操作，n表示所有的像素数，xi和yi分别表示输入的GT_seg和Pred_seg的第i个像素

2.3 loss消融实验

为了强调边界感知精度的重要性，提出了mean Boundary Accuracy measure (mBA) [这与boundary loss很像]，为了对不同大小的图像进行稳健估计，我们在[3，$\frac{w+h}{300}$ ]中采样5个半径,从GT计算每个半径内的分割精度，然后对这些值进行平均。表1显示，我们的模型在IoU方面提高是有效的，在边界精度方面提高更显著。

通过多层级级联，该模块可以将不同的细化阶段委托给不同的尺度。如图3所示，3层模型使用了中间小尺度的分割（将在第3.2节中详细介绍）来更好地捕获对象结构。尽管这些模型都有相同的接受域，3层模型可以更好地利用结构线索，比1层模型产生更详细的分割。

3、Seg结果Refinement

在测试阶段，我们使用Global step和Local step，通过使用相同的训练细化模块来执行高分辨率分割细化。具体来说，Global step考虑整个调整后的图像来修复结构，而Local step使用图像作物以全分辨率细化细节。相同的细化模块可以递归地用于更高分辨率的细化。

3.1 Global step

图4详细介绍了全局步骤的设计，该步骤使用3级cascade重新调整了整个图像。由于测试过程中的全分辨率图像往往不能融入GPU进行处理，我们对输入进行降采样，使长轴的长度为L，同时保持相同的长径比。

对于cascade的输入为其他模型分割好的seg结果，复制输入的seg数据3次，以保持输入信道维数不变。在cascade的第一级之后，其中一个输入通道将被双上采样粗输出取代。这一直重复到最后一个级别，其中输入包括初始分割和来自以前级别的所有输出。
这种设计使我们的网络能够逐步修复分割错误，同时保持初始分割的细节。对于多个层次，我们可以粗略地描述对象，并在粗层次中修复较大的误差，并使用粗层次提供的更鲁棒的特征来关注精细层次中的边界精度。

3.2 Local Step

图5说明了Local Step的详细步骤。由于内存的限制，即使是使用现代gpu，也不能一次性处理非常高分辨率的图像。此外，训练数据和测试数据之间规模的剧烈变化也会导致分割质量较差。我们利用cascade模型，首先使用降采样图像来执行全局细化，然后使用来自更高分辨率图像的图像crop来执行局部细化。这些crop使local setp能够处理没有高分辨率训练数据的高分辨率图像，同时由于全局步骤而考虑到图像上下文。

在Local Step中，模型取Global step的最后一级的两个输出，分别记为$S_4^1$和$S_1^1$。两个输出都被预先调整到图像W的原始大小×h。模型采用L×L的图像crop，从crop输出的两侧裁剪16像素，以避免边界伪影，图像边界除外。crop size以(L/2−32)为变化步幅，这样大多数像素被四种crop size覆盖，超过图像边界的无效crop被移动以与图像的最后一行/列对齐。然后将图像crop输入一个2级的cascade，输出步幅分别为4和1。在融合过程中，来自不同patch的输出可能会由于不同的图像上下文而彼此不一致，我们通过平均所有的输出值来解决这个问题。对于具有更高分辨率的图像，我们可以以从粗到细的方式递归地应用局部步骤。

3.3 L的选择

图6显示了在选择不同的L时，测试期间GPU内存使用和细化质量（mBA）之间的关系。在我们的实验中，我们选择了具有3.16GBGPU内存使用量的L = 900，以平衡增加GPU内存使用量和降低性能增益之间的权衡。在我们对大验证集的实验中，使用更高的L是不必要的，并且占用了额外的内存。在低内存设置中，一个较小的L，如500，可以用来产生稍差的细化（−0.6%mBA），而更低的内存占用（1.16 GB）。请注意，GPU内存的使用只与L有关，而不是图像分辨率，因为融合步骤可以很容易地在CPU上执行。

3.4 Global and Local Refinement的消融实验

表2显示，Global Step和Local Step对于高分辨率分割的细化都是必不可少的。需要注意的是，当我们删除Global Step时，IoU的下降更显著，这表明Global Step主要负责修复整体结构，对IoU的提升贡献更大，而Local Step由于图像上下文不足而无法单独实现。在没有Local Step的情况下，虽然IoU只略有下降，但我们注意到，由于Global Step不能提取高分辨率的细节，边界精度下降得更显著

图7研究了Local Step对于不同分辨率输入的重要性：我们通过调整大验证集的大小来评估使用和没有Local Step的方法。虽然Global Step对于低分辨率的输入是足够的，但Local Step对于尺寸max(W,H)>900的精确高分辨率细化是至关重要的。因此，我们对最大（H，W）≥900的输入同时使用Global Step和Local Step，对较低分辨率的输入只使用Global Step。

实验细节

4.1 训练数据处理

为了学习客观性信息，我们以一种类不可知的方式在一个数据集集合上训练我们的模型。我们合并了MSRA-10K [8]、DUT-OMRON [48]、ECSSD [38]和FSS-1000 [44]，生成一个36,572的分割数据集，比普通数据集如PASCAL（20类）或COCO（80类）更多样的语义类。使用这个数据集（> 1000类）可以使我们的模型更加健壮，并可推广到新的类。

在训练过程中，我们随机选取224x224的image crop，并通过扰动GT-seg值来模拟其他模型的pred-seg结果。输入通过3级级联，在每个级别中计算损失。虽然crop size小于测试中使用的L，但我们的模型设计有助于弥补这一差距。全卷积特征提取器提供了平移不变性，而金字塔池化模块提供了重要的图像上下文，允许我们的模型扩展到更高的分辨率，而没有显著的性能损失。使用较小的crop size加速了我们的训练过程，并使数据准备更加容易，因为用于分割的高分辨率训练数据的获取成本非常昂贵。

为了便于推广，我们避免使用现有模型生成的分割输出进行训练[仅使用扰动GT-seg生成标签]，这可能会导致对特定模型的过拟合。相反，扰动的GT-seg应该描述由其他方法产生的不准确分割的各种形状和输出，这有助于我们的算法对不同的初始分割具有更强的鲁棒性。我们通过对轮廓进行子采样，然后进行随机扩张和侵蚀来产生这种扰动分割。这种扰动的例子如图8所示。

4.2 模型评估数据

我们用PASCAL VOC 2012 [13]、BIG（我们的高分辨率数据集）和ADE20K [55]定量评估了我们的结果。我们在各种设置中不进行任何微调来评估我们的模型，并展示了我们的模型所做的改进。虽然PASCAL VOC 2012数据集被广泛应用于图像分割任务，但它并没有准确的标注图形的边界，边界附近的区域被标记为“空白”。为了进行更准确的评估，我们从PASCAL VOC 2012验证集中重新标记了500个片段，以便在空隙边界区域内找到准确的边界。图9显示了一个重新标记的示例。

BIG数据集： 一个具有50个验证对象和100个测试对象的高分辨率语义分割数据集。图像分辨率的大范围从2048×1600到5000×3600。数据集中的每一张图像都被专业人员仔细标记，同时保持了与PASCAL VOC 2012相同的标注原则，没有空白区域。重新标记的PASCAL VOC验证集和大数据集都可以在我们的项目网站上找到。在评估中使用的其他数据集不被修改。我们使用标准分割度量IoU和边界度量mBA来评估我们的方法。

4.3 实验细节

使用以ResNet为backbone的PSPNet-50[53]作为我们的基础网络。数据增强，包括对GT-seg的扰动，图像翻转和裁剪，都是动态完成的，以进一步增加数据的多样性。我们使用Adam优化器[19]，权重衰减为10−4，对于30K迭代的学习速率为0.0003，然后对于另外30K迭代的学习速率为0.00003，batchsize大小为9。使用2个1080Ti训练，总训练时间约为16小时。Local step只在感兴趣的区域执行，对于图1，完整的细化过程大约需要6.6秒。除非另有说明，我们对所有的实验都使用相同的训练模型。

我们的方法可以仅使用客观性信息来细化输入分割。请注意，我们的模型在训练中从未见过以下任何一个数据集。在本节中，我们将重点关注单个对象的细化效果，我们比较和评估了我们的细化模型对[6,53]在PASCAL VOC 2012数据集上训练的各种语义分割模型输出的影响。我们的方法比常用的多尺度测试更有效，实验结果在补充材料中进一步显示。

4.4 语义分割效果

PASCAL VOC 2012
由于输入模型是在PASCAL VOC 2012数据集中进行训练的，因此不需要调整大小来获得它们的输出，然后将其输入到我们的细化模型中。这些图像的分辨率很低，所以我们可以直接使用全局步骤来重新细化它们。我们在表3的上半部分报告了总体不分类的IoU和边界精度。结果表明，该方法可以在所有情况下提高分割质量，特别是沿边界区域。

BIG
由于内存的限制，大多数现有的分割方法不能直接在全分辨率的大数据集上进行评估。因此，我们通过将调整后的图像输入到现有的模型中来获得初始分割。我们对输入图像进行降采样，使长轴为512像素，并对输出分割进行双次上采样到原始分辨率。在表3的下半部分，我们展示了我们在高分辨率分割的大测试集上的结果。请注意，即使我们从未见过任何高分辨率的训练图像，我们也能够在这些尺度上产生高质量的改进。图14显示了我们的改进的视觉改进。虽然超分辨率模型[11,43]似乎适用生成高分辨率mask，但具有错误分割的输入（例如，图13中缺失的桌子腿和图14中婴儿的手）不能通过超分辨率进行纠正。

我们的方法依赖于输入分割和低水平线索，没有特定的语义能力。图11显示了一个失败的情况，其中输入错误太大，我们的方法无法消除。

4.5 全局分割效果

为了将CascadePSP扩展到场景解析，在类竞争可能存在问题的密集类中，我们提出了一种分治方法，使用我们的预训练网络独立地细化每个语义对象，然后使用融合函数集成结果。图12概述了我们的策略。

我们通过使用具有25%填充的roi，为每个语义对象独立地细化足够大的连接组件。为了处理重叠区域，原始方法将是在输出置信度上使用argmax，这将导致所有类得分较低区域有噪声结果。故此，我们的融合函数是一个修改的argmax，如果所有的输入类置信度值低于0.5，我们采用原始分割结果进行替代。

在这里，我们在ADE20K [55]的验证集上评估了我们的模型。由于ADE20K数据集包含了“stuff”背景类（见补充材料），这些背景类的客观性不强，而且与我们的训练数据差异太大，因此我们降低了它们的输出分数，以专注于前景细化。请注意，细化前景对象仍然可以帮助进行背景细化，因为argmax操作考虑了这两个置信度分数。表4列出了结果，表明我们的模型产生了更高质量的分割。图13显示了样本的定性评价。

5、代码分析

5.1 安装使用

安装命令为 pip install segmentation-refinement
使用示例代码如下，可以通过调整L的值来获取不同的结果，对于大分辨率的影像可以使用fast=True来提升效果。

import cv2
import time
import matplotlib.pyplot as plt
import segmentation_refinement as refine
image = cv2.imread('test/aeroplane.jpg')
mask = cv2.imread('test/aeroplane.png', cv2.IMREAD_GRAYSCALE)

# model_path can also be specified here
# This step takes some time to load the model
refiner = refine.Refiner(device='cuda:0') # device can also be 'cpu'

# Fast - Global step only.
# Smaller L -> Less memory usage; faster in fast mode.
output = refiner.refine(image, mask, fast=False, L=900) 

# this line to save output
cv2.imwrite('output.png', output)

plt.imshow(output)
plt.show()

5.2 CascadePSP代码

CascadePSP的关键代码为RefinementModule，代码地址为 https://github.com/hkchengrex/CascadePSP/blob/master/segmentation-refinement/segmentation_refinement/models/psp/pspnet.py
以下代码描述的forward函数与下图描述的结构完全一致，与官方源码相比添加了具体的代码注释。

class RefinementModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.feats = extractors.resnet50()
        self.psp = PSPModule(2048, 1024, (1, 2, 3, 6))

        self.up_1 = PSPUpsample(1024, 1024+256, 512)
        self.up_2 = PSPUpsample(512, 512+64, 256)
        self.up_3 = PSPUpsample(256, 256+3, 32)
        
        #用于输出步幅为8时的语义分割结果
        self.final_28 = nn.Sequential(
            nn.Conv2d(1024, 32, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 1, kernel_size=1),
        )

        #用于输出步幅为4时的语义分割结果
        self.final_56 = nn.Sequential(
            nn.Conv2d(512, 32, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 1, kernel_size=1),
        )

        #用于融合输出步幅为1时的特征图 concat(final_56(x),x)
        self.final_11 = nn.Conv2d(32+3, 32, kernel_size=1)
        #用于输出步幅为1时的语义分割结果
        self.final_21 = nn.Conv2d(32, 1, kernel_size=1)

    def forward(self, x, seg, inter_s8=None, inter_s4=None):

        images = {
    
    }

        """
        First iteration, s8 output
        """
        if inter_s8 is None:
            p = torch.cat((x, seg, seg, seg), 1)

            f, f_1, f_2 = self.feats(p) 
            p = self.psp(f)

            inter_s8 = self.final_28(p)
            r_inter_s8 = F.interpolate(inter_s8, scale_factor=8, mode='bilinear', align_corners=False)
            r_inter_tanh_s8 = torch.tanh(r_inter_s8)#用于做s4时第二次级联时的输入

            images['pred_28'] = torch.sigmoid(r_inter_s8) #用于做输出的结果
            images['out_28'] = r_inter_s8 #无用变量
        else:
            r_inter_tanh_s8 = inter_s8

        """
        Second iteration, s8 output
        """
        if inter_s4 is None:
            p = torch.cat((x, seg, r_inter_tanh_s8, r_inter_tanh_s8), 1)

            f, f_1, f_2 = self.feats(p) 
            p = self.psp(f)
            inter_s8_2 = self.final_28(p)
            r_inter_s8_2 = F.interpolate(inter_s8_2, scale_factor=8, mode='bilinear', align_corners=False)
            r_inter_tanh_s8_2 = torch.tanh(r_inter_s8_2)#重新生成第一次级联时的seg
            p = self.up_1(p, f_2)

            inter_s4 = self.final_56(p)
            r_inter_s4 = F.interpolate(inter_s4, scale_factor=4, mode='bilinear', align_corners=False)
            r_inter_tanh_s4 = torch.tanh(r_inter_s4)#用于做s1时第三次级联时的输入

            images['pred_28_2'] = torch.sigmoid(r_inter_s8_2)
            images['out_28_2'] = r_inter_s8_2
            images['pred_56'] = torch.sigmoid(r_inter_s4)
            images['out_56'] = r_inter_s4
        else:
            r_inter_tanh_s8_2 = inter_s8
            r_inter_tanh_s4 = inter_s4

        """
        Third iteration, s1 output
        """
        p = torch.cat((x, seg, r_inter_tanh_s8_2, r_inter_tanh_s4), 1)

        f, f_1, f_2 = self.feats(p) 
        p = self.psp(f)
        inter_s8_3 = self.final_28(p)
        r_inter_s8_3 = F.interpolate(inter_s8_3, scale_factor=8, mode='bilinear', align_corners=False)

        p = self.up_1(p, f_2)
        inter_s4_2 = self.final_56(p)
        r_inter_s4_2 = F.interpolate(inter_s4_2, scale_factor=4, mode='bilinear', align_corners=False)
        p = self.up_2(p, f_1)
        p = self.up_3(p, x)


        """
        Final output
        """
        p = F.relu(self.final_11(torch.cat([p, x], 1)), inplace=True)
        p = self.final_21(p)

        pred_224 = torch.sigmoid(p)

        images['pred_224'] = pred_224
        images['out_224'] = p
        images['pred_28_3'] = torch.sigmoid(r_inter_s8_3)
        images['pred_56_2'] = torch.sigmoid(r_inter_s4_2)
        images['out_28_3'] = r_inter_s8_3
        images['out_56_2'] = r_inter_s4_2

        return images