图像分割-RSPrompter

文章目录

前言

《RSPrompter: Learning to prompt for remote sensing instance segmentation based on visual foundation model》，2024

本篇论文提出了目前SAM存在的一些问题：

SAM严重依赖于人工先验（点、框、mask）且分割结果是没有标记类别的，如下图，单点、两点、边框提示的不同分割结果
SAM在遥感图像分割任务中的性能在很大程度上仍未被探索和证明

1. 自动化提示器

针对问题1，不由得想到：如果能够自动生成多个与类别相关的提示，SAM 的解码器就能够产生带有类别标签的多个实例级掩码。

论文提出训练一个自动化的提示器（RSPrompter），能够处理来自测试集的任何图像，同时对对象进行定位，并推断它们的语义类别和实例掩码。最后通过RSPrompter+SAM，实现分割任务的自动化。

这个过程存在两个主要挑战：（i）类别相关的提示从哪里来？（ii）应选择哪种类型的提示作为掩膜解码器的输入

RSPrompter的结算流程如下：

通过冻结权重的SAM image encoder生成多个中间特征图 $F_i$ （K×h×w×c）和最终的图像嵌入 $F_{img}$ （h×w×c）
$F_i$ 通过多尺度特征增强器（结构和原理见1.1）逐步处理获得多尺度特征 $F_{ms}$
将多尺度特征 $F_{ms}$ 输入RSPrompter（结构和原理见1.2）以获得多组提示语（ $F^m_{sparse} ∈ R^{K_p×c}, m ∈ \{1, · · · , Np\}$ ）和它们的语义类别（ $c_m ∈ R^c, m ∈ \{1, · · · , Np\}$ ），其中， $K_p$ 定义了每次掩码生成的提示嵌入次数； $N_p$ 是提示的个数，用来定义输出实例掩码的个数。

需要注意的是， $F^m_{sparse}$ 只包含前景目标实例提示，其语义类别由 $c_m$ 给出。单个 $F^m_{sparse}$ 是多个提示的组合，即用多个点嵌入或一个边框表示一个实例掩码。

1.1 多尺度特征增强器

为了在不增加提示器（Prompter）计算复杂度的情况下提取具有语义相关性和判别性的特征，论文提出了一种轻量级的多尺度特征增强器。该增强器包括特征聚合器（Feature Aggregator）和特征分割器（Feature Splitter）。

Feature Aggregator的设计是为了从SMA ViT backbone的众多中间特征层中学习具有代表性的语义特征并进行融合，Feature Splitter被用于从融合后的特征图中生成多尺度金字塔特征图。具体过程如下：

对每个中间层特征图下采样（1×1卷积降维+3×3卷积提升空间信息）到同一通道：h×w×c -> h×w×32
通过跳连+add形势逐渐合并各层特征，FusionConv表示最终的融合卷积层，由两个3 × 3卷积层和一个1 × 1卷积层组成，以恢复通道维度，得到最终融合特征图 $F_{agg}$ 。
将 $F_{agg}$ 输入特征分割器，Feature Splitter使用转置卷积层生成上采样特征，使用最大池化生成下采样特征。通过利用上/下采样层，最终获得了五种不同尺度的特征图： $F_{ms}^j ∈ R^{\frac{H}{2^{j+1}},\frac{W}{2^{j+1}},c}$ ，其中 $j ∈ \{1, 2, 3, 4, 5\}$ 。

1.2 RSPrompter

在获取语义特征后，利用提示器为SAM mask decoder生成提示嵌入就变得可行。采用了两种不同结构的提示器：anchor-based和query-based。

Anchor-based Prompter

首先利用基于锚点的区域建议网络( RPN )生成候选目标框。随后通过RoI Pooling从位置编码的特征图中提取每个对象的独特视觉特征表示。

然后利用对象的视觉特征得到3个感知头：语义头（semantic head）、定位头（localization head）和提示头（prompt head）。

语义头部的作用是识别特定的物体类别
定位头部负责建立生成的提示表示与目标实例掩码之间的匹配准则，即基于定位( Intersection over Union , IoU)的贪婪匹配。
提示头为SAM mask decoder生成必要的提示嵌入。

在这里插入图片描述

ps：PE 是positional encoding

简单来说，该提示器本质上就是一个简单的目标检测网络。其主要框架和Faster RCNN基本一致。其损失函数如下：

在这里插入图片描述

包含：

候选区生成网络损失 $L_{rpn}$
分类损失 $L_{cls}$ ：交叉熵(CE)
回归损失 $L_{reg}$ ：SmoothL1 loss for 偏移量
分割损失 $L_{seg}$ ：表示SAM decoder 掩码与真实实例掩码之间的二进制CE损失，其中框的IoU决定了有监督的匹配准则。

Query-based Prompter

Anchor-based的方法较为复杂，需要使用框信息进行掩码匹配和监督训练。

因此提出了以最优传输为基础的query-based提示器。如下图所示：

编码器用于提取高层次聚集的语义特征，建模各尺度特征图间的关系，输出建模后的多尺度特征图 $\hat F_i$ ，其中 $\hat F_m$ 是尺寸最大的特征图（原图四分之一大小）
解码器通过交叉注意力交互将预设的可学习查询token转换为SAM和相应语义类别的提示嵌入
- 可学习的查询token： $F^i_{query} ∈ R^{N_p×c}$ ，其中 $N_p$ 为提示个数，即实例数量
- 可学习的查询token和多尺度特征图进行交叉注意力计算再经过三个head后得到：
  - 输出类别： $\hat c_i ∈ R^{N_p×c}$
  - 输出掩膜滤波器： $\hat f_i ∈ R^{N_p×c}$
  - 输出提示嵌入 $\hat c_i ∈ R^{N_p×K_p×c}$ ，其中 $K_p$ 用于定义每个提示的嵌入次数，即表示一个实例目标所需的提示数量。
- 第i级粗分割掩膜（ $\hat m_{coarse}^i$ ）由 $\hat F_m$ 通过 $\hat f_i$ 线性加权得到
- 然后将 $\hat m_{coarse}^i$ 经过SAM promopt encoder编码后得到稠密提示编码 $F^i_{dense}$
- 通过计算 $e_i + sin(e_i)$ 后得到稀疏提示编码 $F^i_{sparse}$

整体方程可以沿着层级i循环计算，以获得多语义结果。为了进行推理，只考虑最后一层，通过数学运算得到二值掩码的边框。

在这里插入图片描述

在基于查询的提示器的训练过程中，主要进行了两个步骤：( i ) $N_p$ 预测掩码与K个真实实例掩码(一般来说, $Np > K$ )的匹配；( ii )利用匹配的标签进行监督训练的后续实现。

1. 在执行最优运输匹配时，建立如下的匹配成本，该成本同时包含了预测的类别和掩码：

在这里插入图片描述

ω 表示分配关系。使用匈牙利算法来确定 $N_p$ 个预测和K个目标之间的最优分配。匹配代价考虑了预测和真实标注之间的相似性。具体来说，它包括类分类匹配代价( $L_{cls}$ )、掩码交叉熵代价( $L_{seg-ce}$ )和掩码dice代价（ $L_{seg-dice}$ 是一种相似性度量损失，和iou loss类似）。

2. 匹配后即可计算最终loss，由分类交叉熵损失和mask二元交叉熵损失构成：

在这里插入图片描述

2. SAM的扩展

针对上述问题2：除了提出的RSPrompter，还引入了另外三种基于SAM的实例分割方法进行对比分析，以探索SAM性能改进方向。其中，前两种方法为论文提出，具有简单实现的SAM - det已经在社区内引起了相当大的关注和应用。

在这里插入图片描述

（a）SAM-seg：冻结SAM Image Encoder，使用多尺度增强器获取多尺度特征，然后输入到其他网络的实例分割头，即Mask RCNN或Mask2Former的mask decoder。
（b）SAM-cls：在图像上均匀分布点，并将每个点作为实例的提示输入，将提示输入到SAM以获得图像中所有潜在的实例目标。在获得所有实例的掩码后，再使用分类器为每个掩码分配标签。
（c）SAM-det：首先，训练一个目标检测器，以精确定位图像中的期望目标。随后，将检测到的边界框作为提示输入到SAM中。
（d）SAM-RSPromopt：和SAM-det的理念是比较相似的，区别在于直接生成提示嵌入，而非提示本身。