SAM2-UNet: SAM2用于自然和医学图像分割

SAM2-UNet: Segment Anything 2 Makes Strong Encoder for Natural and Medical Image Segmentation

https://arxiv.org/pdf/2408.08870

Project page

Abstract

Image segmentation plays an important role in vision understanding. Recently, the emerging vision foundation models continuously achieved superior performance on various tasks. Following such success, in this paper, we prove that the Segment Anything Model 2 (SAM2) can be a strong encoder for U-shaped segmentation models. We propose a simple but effective framework, termed SAM2-UNet, for versatile image segmentation. Specifically, SAM2-UNet adopts the Hiera backbone of SAM2 as the encoder, while the decoder uses the classic U-shaped design. Additionally, adapters are inserted into the encoder to allow parameter-efficient fine-tuning. Preliminary experiments on various downstream tasks, such as camouflaged object detection, salient object detection, marine animal segmentation, mirror detection, and polyp segmentation, demonstrate that our SAM2-UNet can simply beat existing specialized state-of-the-art methods without bells and whistles.

本文证明了 Segment Anything Model 2 (SAM2) 对于U-Net图像分割任务非常有效。

本文提出了图像分割的 SAM2-UNet。

SAM2-UNet 使用 SAM2 的分级骨干网络作为编码器；解码器使用经典U-Net。

微调过程的adapters插在编码器中。

Method

The overall architecture of SAM2-UNet is illustrated in Fig. 1, comprising four main components: encoder, decoder, receptive field blocks (RFBs), and adapters. Note that we discard components that are not essential for constructing a basic U-Net [37], such as memory attention, prompt encoder, memory encoder, and memory bank.

编码器。SAM2-UNet采用了由SAM2预训练的Hiera骨干网络。与SAM1中使用的普通ViT编码器相比，Hiera采用了一种层次结构，允许捕获多尺度特征，这更适合于设计U形网络。对于Hiera-L，Ci的取值分别为{144, 288, 576, 1152}。

RFBs。在提取编码器特征后，通过四个RFBs来减少通道数至64，并增强这些轻量级特征。

适配器。由于Hiera的参数可能非常庞大（Hiera-L为2.14亿），因此进行完全微调在内存上可能并不总是可行的。因此，本文冻结Hiera的参数，并在Hiera的每个多尺度块之前插入适配器，以实现参数高效的微调。每个适配器由一个用于下采样的线性层、一个GeLU激活函数、另一个用于上采样的线性层以及最终的GeLU激活函数组成。

解码器。SAM2中的原始掩码解码器采用双向转换器方法，以促进提示嵌入和编码器特征之间的特征交互。本文的解码器使用经典的U-Net设计。它由三个解码器块组成，每个块包含两个“Conv-BN-ReLU”组合，其中“Conv”表示3×3卷积层，“BN”表示批归一化。每个解码器块的输出特征通过1×1卷积分割头来生成分割结果Si（i ∈ {1, 2, 3}），然后对其进行上采样，并由真实掩码G进行监督。

损失函数。使用加权IoU损失和二元交叉熵（BCE）损失作为训练目标：

。

此外，对所有分割输出Si应用深度监督。SAM2-UNet的总损失公式为：

。

Dataset

伪装目标检测（Camouflaged Object Detection）旨在检测环境中隐藏良好的物体。我们采用四个数据集进行基准测试，包括CAMO [19]、COD10K [8]、CHAMELEON [40]和NC4K [30]。用于比较的四项指标包括S-measure（Sα）[6]、自适应F-measure（Fβ）[31]、平均E-measure（Eϕ）[7]和平均绝对误差（MAE）。

显著目标检测（Salient Object Detection）旨在模仿人类认知机制来识别显著物体。我们采用五个数据集进行基准测试，包括DUTS [44]、DUT-O [52]、HKU-IS [20]、PASCAL-S [24]和ECSSD [51]。用于比较的三项指标包括S-measure（Sα）[6]、平均E-measure（Eϕ）[7]和平均绝对误差（MAE）。

海洋动物分割（Marine Animal Segmentation）专注于探索水下环境以发现海洋动物。我们采用两个数据集进行基准测试，包括MAS3K [21]和RMAS [10]。用于比较的五项指标包括mIoU、S-measure（Sα）[6]、加权F-measure（F w β）[31]、平均E-measure（Eϕ）[7]和平均绝对误差（MAE）。

镜面检测（Mirror Detection）能够识别给定输入图像中的镜面区域。我们采用两个数据集进行基准测试，包括MSD [53]和PMD [25]。用于比较的三项指标包括IoU、F-measure [31]和平均绝对误差（MAE）。

息肉分割（Polyp Segmentation）有助于结直肠癌的诊断。我们采用五个数据集进行基准测试，包括Kvasir-SEG [17]、CVC-ClincDB [1]、CVC-ColonDB [42]、CVC-300 [43]和ETIS [39]。用于比较的两项指标包括平均Dice（mDice）和平均IoU（mIoU）。

Results