点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
随着大规模建模的最新突破,Segment Anything Model(SAM)在各种视觉应用中展现出巨大潜力。然而,由于缺乏水下领域专业知识,SAM及其变体在端到端水下实例分割任务中面临性能限制,而它们更高的计算需求进一步阻碍了它们在水下场景中的应用。
为应对这一挑战,作者提出了大规模水下实例分割数据集UIIS10K,其中包含10,048张带有像素级标注的10类图像。随后,作者引入了UWSAM,这是一种专为水下实例自动和精确分割而设计的高效模型。
UWSAM通过基于Mask GAT的水下知识蒸馏(MG-UKD)方法,将SAM ViT-Huge图像编码器的知识高效地蒸馏到更小的ViT-Small图像编码器中,以实现有效的视觉表征学习。
此外,作者为UWSAM设计了端到端水下 Prompt 生成器(EUPG),该生成器自动生成水下 Prompt ,而不是显式地提供前景点或框作为 Prompt ,从而使网络能够准确定位水下实例,实现高效分割。
综合实验结果表明,yinqing-UWSAM_2505是有效的,在多个水下实例数据集上显著优于现有方法。
水下视觉理解与密集预测任务近年来在计算机视觉领域受到越来越多的关注,旨在探索和保护水下生态系统[1]-[5]。实例分割是一项基础且至关重要的视觉任务,旨在分割图像中的所有物体并相互区分,从而实现详细准确的场景理解[6]。通过提供像素级物体定位,实例分割为更精细的实例级感知和广泛应用奠定了基础[7]。水下图像实例分割(UIIS)在水下视觉应用扩展方面展现出巨大潜力,例如海洋遗迹发现、海洋资源探索、水下人机交互以及水下图像综合分析[4]。
近年来,大语言模型(LLMs)如生成式预训练 Transformer (GPT-4)[8]、自适应多任务架构语言学习(LLaMA)[9]和路径语言模型(PaLM)[10]在自然语言处理(NLP)领域引发了革命。这些基础模型展现出优异的泛化能力,并在众多开放世界语言任务中表现优异。受LLMs成功的启发,视觉基础模型如对比语言图像预训练模型(CLIP)[11]、任意分割模型(SAM)[12]和任意分割模型2(SAM2)[13]也相继出现。这些基础模型的引入持续推动着计算机视觉领域的研究者探索。值得注意的是,SAM和SAM2凭借其强大的编码器-解码器 Transformer 框架和大规模数据集,在各种分割任务中表现出色。通过微调或适当修改,它们在海洋科学领域具有巨大潜力。尽管具有这些优势,将SAM等模型部署在无人机载等边缘计算设备上仍面临性能 Bottleneck 。主要问题源于模型架构的复杂性,特别是图像编码器(如ViT-H [14]和Hiera-L [10]),分别包含636M和224M参数。因此,使用这些模型进行实例分割需要大量的计算和内存成本,极大地限制了它们在水下环境中的应用。
为应对这一挑战,近期多项研究提出了避免在基于 Prompt 的实例分割中产生高昂成本的策略。例如,FastSAM [15] 采用基于实时CNN的架构以降低计算开销,而EfficientSAM [16] 则利用著名的 Mask 自编码器(MAE)方法 [17] 和SAM模型来优化轻量级且高效的图像编码器。
尽管付出了这些努力,但由于缺乏专门的水下视觉知识,SAM及其变体在水下实例分割任务中仍未取得令人满意的表现。为解决这一局限性,作者提出了UIIS10K数据集,这是一个包含10,048张图像的大型综合性水下实例分割数据集,其中包含鱼类、珊瑚、遗迹和人类等10个类别的像素级标注。作者的数据集旨在为水下实例分割任务提供坚实的基础,促进开发更适合水下环境的模型。
同时,作者引入了UWSAM,这是一种基于SAM框架设计用于水下图像内在特征的多目标水下图像实例分割方法。UWSAM采用基于Mask GAT的水下知识蒸馏(MG-UKD)算法来训练ViTSmall Backbone 网络[14]作为图像编码器,而端到端水下 Prompt 生成器(EUPG)与SAM Mask解码器结合引导网络实现端到端分割。MG-UKD利用了水下图像的一个基本特征:水下实例(如鱼群、珊瑚礁)通常形成群体和集群,导致水下图像的不同区域具有视觉相似性。利用这一先验知识,作者随机 Mask ViT-Small图像编码器输出的特征,并采用图注意力网络(GAT)来重建这些特征,使其与SAM ViT-Huge图像编码器提取的特征相匹配。通过将知识从更大、更复杂的模型蒸馏到更小、更高效的模型中,得到的模型更适合水下部署。通过结合水下领域特定知识并采用有效的特征重建,yinqing-UWSAM_2505成功解决了水下环境的独特挑战,包括光照条件变化、浊度和多种海洋生物的存在。
此外,与传统方法不同,后者依赖于外部目标检测器生成用于SAM的点或边界框(BBox) Prompt ,随后由SAM Prompt 编码器将其编码为 Prompt 特征,EUPG直接将水下 Prompt 特征的生成集成到模型内部,从而实现端到端训练和推理。通过引入EUPG,yinqing-UWSAM_2505消除了对手动提供的 Prompt 或外部检测器的需求,显著提高了水下分割的效率和适应性。实验结果表明,EUPG不仅提高了分割精度,还降低了计算复杂度,使其非常适合部署在资源受限的水下平台。
为评估yinqing-UWSAM_2505的有效性,作者在UIIS10K数据集、USIS10K数据集[4]以及UIIS数据集[2]上进行了广泛的实验,将yinqing-UWSAM_2505与基于SAM的最先进方法以及其他水下和通用实例分割方法进行了比较。结果表明,yinqing-UWSAM_2505在分割精度、鲁棒性和效率方面显著优于这些现有方法。yinqing-UWSAM_2505尤其擅长解决水下环境特有的挑战,例如光照变化、遮挡以及水下物体复杂的表观特征。yinqing-UWSAM_2505卓越的性能突显了其在水下探索和监测中的潜力,即使在具有挑战性的条件下也能提供准确高效的分割能力。主要贡献总结如下:
作者提出了UIIS10K,一个包含10,048张图像的大型水下实例分割数据集,其中10个类别使用UWSAM模型进行了像素级标注。据作者所知,这是目前可用的最大规模的水下实例分割数据集,可用于评估水下分割方法。
作者提出了基于Mask GAT的水下知识蒸馏(MG-UKD)算法,用于UWSAM,该算法将大型SAM ViT-Huge编码器的知识蒸馏到更适合水下环境的ViT-Small编码器中。MG-UKD降低了计算复杂度,并在低能见度和光照变化等挑战性条件下提高了分割精度,使其非常适合实际水下应用。
作者为UWSAM开发了端到端水下 Prompt 生成器(EUPG),该系统能够直接生成包含位置信息和上下文细节的 Prompt ,无需外部检测器或手动输入,从而实现高效的水下端到端分割。
A. 水下实例分割数据集
与陆地分割模型不同,水下实例分割由于水下发生的光学畸变,如光散射、折射和颜色衰减 [18],面临着独特的挑战。这些因素导致图像通常模糊不清、对比度差、模糊或颜色失真,给通用分割模型带来困难。这种退化可以通过SeaThru模型 [19] 进行定量描述:
β
β
在通道
中,
是相机捕获的图像,
是场景辐射,
是 Mask 光,
和
控制器分别是图像在RGB通道中的光衰减系数和后向散射系数 [20], [21],
是场景深度。除了质量退化,海洋生物在形状、大小和颜色上表现出显著多样性 [22], [23],所有这些因素导致在陆地实例分割数据集上训练的通用模型无法在水下实例分割任务中取得令人满意的结果。因此,近年来社区一直专注于标注水下分割数据集。TrashCan数据集 [24] 专注于海洋环境保护,为训练自主水下航行器(AUVs)检测器提供了标注的水下垃圾图像。DeepFish数据集 [25] 包含批发鱼市场的空中鱼图像,用于鱼类物种分类和尺寸估计。UIIS数据集 [2] 是第一个通用型水下实例分割数据集,包含7个类别的4628张图像,具有像素级标注,从而鼓励了更精确和上下文敏感的分割解决方案的开发。USIS10K数据集 [4] 专注于实例分割的另一种变体——水下显著实例分割,旨在帮助分割模型专注于分割有价值的水下显著目标。然而,上述数据集要么是为特定任务设计的,要么标注数量不足。由于在大模型时代,大规模基准数据集在开发水下实例分割方法中起着至关重要的作用,作者构建了一个新的具有挑战性的数据集,称为USIs1OK,涵盖10,048张水下图像,为10类水下实例提供细粒度标注。与其他现有数据集相比,作者的数据集规模更大,并且具有更多样化的物体类别。
B. 实例分割模型
地面实例分割研究在过去十年中取得了显著进展[26]。经典的实例分割方法通常依赖于Mask R-CNN[6]的两阶段框架,该框架使用区域 Proposal 网络(RPN)[27]生成边界框,并使用RoIAlign[6]从特征金字塔[28]中提取实例特征,以进行像素级 Mask 预测。随着Transformer[29]在计算中的应用,QueryInst[30]受DETR[31]启发,将目标视为可学习的分割 Query ,并使用动态 Mask 头进行分割,在准确性和速度上均优于先前技术。此外,Mask2Former[32]通过限制Transformer解码器中的交叉注意力范围来提升分割性能,从而更有效地提取局部特征。虽然水下实例分割的历史与实例分割任务相比较短,但已有相当多的研究工作与此相关。WaterMask[2]使用多尺度精炼和图注意力机制来对抗水下视觉伪影,从而提高水下实例分割的准确性。TC-USOD[5]在Transformer编码器中深度融合RGB图像与深度图像,并使用轻量级卷积解码器分割物体的显著性 Mask 。最近,基础模型(例如CLIP[11]、GLIP[33]、ALIGN[34]、SAM[12]、SAM2[13])在计算机视觉领域受到了广泛关注。其中,SAM[12]是一个在大型SA-1B数据集上训练的交互式分割模型,以语义无关的方式接收各种用户输入(例如点、框和 Mask ),以准确分割视觉目标,展现出强大的零点泛化能力和可 Prompt 的分割性能。CoralsCOP[1]将SAM应用于珊瑚礁分割任务,通过在SAM解码器中添加并行语义分支,并在训练中引入带有"非珊瑚"标签的负 Mask 来帮助网络收敛,从而解决珊瑚礁分类挑战。此外,MarineInst[3]结合SAM与冻结的视觉语言模型(如CLIP[11]或MarineGPT[35]),通过实例级视觉描述分析水下图像,支持广泛的海洋视觉分析和场景理解任务。然而,上述所有模型都依赖于用户输入的外部显式 Prompt ,无法实现端到端的水下实例分割。此外,大量的参数进一步限制了它们在水下车辆等场景中的应用。为此,作者提出MG-UKD算法通过知识蒸馏减少UWSAM的参数数量,并引入EUPG模块生成视觉 Prompt ,以实现自动水下实例分割。
C. 知识蒸馏
知识蒸馏(KD)已成为一种广泛使用的将知识从大型教师模型迁移到小型学生模型的技术,同时保持高性能[36]。Hinton等人[37]的基础性工作引入了使用教师模型的软标签来指导学生模型的思想,提高了学生模型泛化能力,尽管其参数较少。自那时以来,大量研究探索了KD的各个方面,旨在提高其效率和有效性。MGD[38]从学生模型的特征中随机遮盖像素,并允许学生通过简单的卷积操作重建教师模型的完整特征,从而鼓励学生关注数据的重要特征。DMAE[39]专注于从预训练模型(如 Mask 自编码器MAE[17])中提取知识,并通过最小化教师模型和学生模型在可见输入块上提取的中间特征之间的距离来执行高效的知识蒸馏。EfficientSAM[16]也将MAE应用于知识蒸馏,引入了SAMI方法。然而,与DMAE不同,SAMI仅遮盖学生模型的特征,并使用额外的交叉注意力解码器来重建这些特征,以使其与教师模型保持一致。然而,由于 Backbone 模型通常在陆地数据集上进行预训练,缺乏水下视觉知识,直接使用其蒸馏结果进行水下实例分割可能导致次优性能。
为推进Ulls研究的发展,作者收集并标注了一个名为Ulls10K的新大规模基准数据集,该数据集包含总共10,048张RGB水下图像,具有像素级实例标注。该数据集的部分样本展示在图1中。在本节中,作者将详细描述Ulls10K的构建过程,并全面分析其特征。
A. 数据集收集与标注与划分
构建UIls1OK数据集分为以下三个步骤:1)图像收集和过滤,2)图像标注,3)数据集划分。
数据集收集。为了尽可能丰富数据集中的图像种类,覆盖更多水下场景和光照条件,作者从互联网和不同领域公开的水下数据集中收集了约40,000张图像,这些数据集包括水下图像增强[40]-[42]、水下语义分割[43]、水下实例分割[2]、[24]以及水下显著性检测[4]、[5]、[44]等。这些图像来自各种水下环境,包括深海、浅水和湖泊,涵盖海洋资源勘探、智能人机协作和水下环境保护等任务。随后,作者安排两名志愿者筛选候选图像,移除重复、损坏或非水下环境的图像。最终,作者仔细标注了剩余的14,500张图像。
数据集标注。利用SAM强大的零样本泛化能力,显著减少了标注工作量,作者采用SAM [12] 和 EfficientSAM [16] 模型进行辅助标注。模型生成初始标签后,志愿者手动进行细化,并为每个 Mask 分配类别标签。作者招募了16名志愿者对数据集进行标注。在开始标注过程之前,他们接受了以下四个方面的培训:
(1) 水下场景常见生物的分类。
(2) 之前数据集构建过程和基本图像标注方法。
(3) 如何使用SAM和EfficientSAM模型进行半自动标注的有效 Prompt 输入。
(4) 如何在半自动标注后手动细化 Mask 。
作者使用Sparse标注多边形来 Token 数据集中的每个实例。标注数据将存储在广泛使用的COCO风格格式[45]中,以确保与大多数流行框架和模型的兼容性。每张图像将由至少两名志愿者标注,并由第三名志愿者进行审核。对于从水下实例分割数据集中收集的图像,两名志愿者将手动细化原始实例标注并重新分配标签。然后,第三名志愿者将选择、细化和合并他们的标注,并检查图像是否包含未标注的实例。对于来自水下显著分割数据集的图像,除了细化和重新分配现有 Mask 的标签外,两名志愿者还将使用SAM和EfficientSAM标注先前 未标注 的非显著实例。第三名志愿者将选择最佳标注并进行细化。对于其他图像,两名志愿者将使用SAM和EfficientSAM在实例 Level 标注所有感兴趣的目标。第三名志愿者将选择并合并他们的标注,进一步细化 Mask 和类别标签以确保准确性和精确性。在使用模型进行辅助标注时,SAM和EfficientSAM针对志愿者提供的每组 Prompt 各生成三个候选 Mask 。志愿者随后选择最合适的 Mask 并手动细化。
表1 与现有水下实例数据集的比较。其中数字表示每个数据集中的图像数量
表2 UIIS1OK 数据集类别描述
为了提高类别标注的准确性,作者根据[46]、[47]中概述的指南对潜在的模糊目标进行了分类。此外,作者排除了那些在标注上无法达成共识的图像实例。最终,作者获得了10,048张图像,这些图像构成了UIIS10K数据集。
数据集划分。为确保在UIIS10K数据集上对深度学习方法进行实用训练和可靠的测试结果,必须在训练集和测试集中包含每个类别的足够样本。因此,作者遵循UIIS10K数据集训练集和测试集约8:2的划分方式。具体而言,UIIS10K数据集被划分为8,083个样本用于训练,以及2,010个样本用于验证和测试。
B. 数据集特征与统计
在本节中,作者阐述了UIlS10K数据集的基本信息、特征和挑战。
数据集的数量和类别。为了提高网络在复杂海洋环境中的泛化能力并避免过拟合,需要大量来自不同场景且包含不同类别的水下图像。为此,作者扩展了UIIS数据集[2],为10,048张来自不同场景的图像提供了详细的标注。如表1所示,UIIS10K是现有水下场景下最大的实例分割数据集。此外,为了使UIIS10K数据集适用于更广泛的下游任务,作者根据表2中所示的10个类别对UIIS10K中的实例进行了标注。UIS10K数据集包括鱼、爬行动物、偶蹄动物和软体动物等类别,这些是海洋生态探索的主要研究目标。它还提供了珊瑚、植物和垃圾的像素级标注,这些是海洋生态保护的主要目标。在图2(b)中,作者还统计了每个类别的 Mask 数量。UIIS10K数据集总共包含41,862个 Mask ,其中鱼和珊瑚类别的 Mask 数量最多,它们是水下环境中最常见的两个目标。此外,数据集还包含人类潜水员、机器人和废墟的标注,这些对于涉及人-机-物智能协作的训练应用非常有价值。详细的类别定义见表2。此外,UIIS10K数据集将所有标注分为三个部分:类别标签、实例 Mask 和边界框。这为使用UIIS10K数据集进行水下场景下的其他下游任务(如目标检测和语义分割)提供了可能性。
图像中的实例数量。在UIIS10K数据集中,单张图像通常包含多个实例。如图1所示,这些图像中的大多数通常包含超过一个实例。在图2(a)中,作者提供了数据集中实例数量的统计数据。具体而言,对于UIIS10K数据集,22.83%的场景包含超过5个实例,8.45%的场景包含超过10个实例,包含最多实例的图像有81个实例。水下图像中实例数量的增加通常是由于鱼群或珊瑚群等因素。例如,在图1第二行第二列的图像中,实例倾向于紧密聚集并相互遮挡,这给准确分割实例边界带来了更大的挑战。
水下图像的通道强度。在水中环境捕获的光学图像由于水对不同波长的选择性吸收,不可避免地表现出颜色衰减,其中红色通道的衰减程度比蓝色和绿色通道高一个数量级[48]。因此,水下图像通常呈现蓝色或绿色。为了量化UIIS10K的这种衰减特性,作者计算了每幅图像中R、G和B通道的平均通道强度和概率密度,如图3所示。结果表明,红色通道的强度最低,但其趋势与绿色和蓝色通道相似。
实例大小。分割过小或过大的实例是实例分割领域中的一个常见但具有挑战性的问题。在[45]中,实例被分为三个 Level :小型(面积 < 32²)、中型(32² ≤ 面积 < 96²)和大型(面积 ≥ 96²),其中面积表示实例中的像素数量。在UIIS10K数据集中,这三个 Level 的实例数量分别为15,800、13,702和12,360,比例为1.27:1.05:1。这导致UIIS10K中实例大小分布近乎均衡,有助于网络有效学习分割不同大小的实例。多种图像分辨率和图像场景。如图2(c)所示,UIIs10K数据集包含各种分辨率的图像,从手持相机拍摄的低分辨率图像到工业设备在水下任务中拍摄的中等分辨率图像,以及少量高分辨率图像。这种多样的分辨率范围旨在满足不同任务的需求。此外,UIIS10K数据集还包含浅水图像(例如图1第6列第2行)、质量显著下降的图像(例如图1第4列第2行)、背景复杂的图像(例如图1第2列第2行)以及饱和度或对比度高的图像(例如图1第4列第1行)。通过这些场景,UIIS10K可以全面评估网络在不同水下场景中的泛化能力。
tokens
(包括 Mask token和IoU token)。SAM的推理过程可以表示如下:
A. 初步
首先,作者概述了SAM架构。SAM由三个主要组件构成:图像编码器
、 Prompt 编码器
和 Mask 解码器
。图像编码器
利用基于 Mask 自编码器(MAE)方法训练的视觉Transformer(ViT)[17]。具体而言,SAM通常采用ViT-Huge变体。 Prompt 编码器
能够处理外部提供的Sparse Prompt (例如点、框)
和密集 Prompt (例如 Mask )
,并将它们转换为 Prompt 特征。 Mask 解码器
基于改进的Transformer解码块[29]。通过双向交叉注意力机制, Mask 解码器
对Sparse Prompt 特征
和图像特征
之间的交互进行建模,并通过预插入的可学习模块预测输出 Mask
及相应的IoU分数
,其中
表示输入图像,
表示特征组合。SAM2通过引入记忆注意力机制和记忆库等组件将SAM扩展到视频分割,但由于作者的关注点在于图像分割,且SAM2在此情况下与SAM基本采用相同结构,仅进行了如将Hiera-Large[49]用作图像编码器等微小修改,因此作者不再详细阐述。
C. 基于 Mask GAT的水下知识蒸馏
基于大规模通用图像分割数据集SA-1B训练的SAM图像编码器,具备强大的图像特征提取能力和丰富的视觉知识,对于知识蒸馏学生模型具有重要价值。然而,由于SA-1B数据集与复杂海洋环境之间存在领域偏差,SAM图像编码器缺乏水下环境的领域专业知识。此外,现有的知识蒸馏方法并非专门针对水下环境设计,在蒸馏过程中未融入水下先验知识。因此,直接使用SAM图像编码器作为教师模型来训练水下实例分割模型可能导致性能欠佳。为解决上述问题,作者提出了MG-UKD。MG-UKD的知识蒸馏过程借鉴了水下图像的一个特征:在海洋环境中,物体常以集群形式出现(如鱼群、珊瑚群),导致相似视觉信息在不同图像区域重复出现。因此,与MGD [38]或SAMI [16]相比,MG-UKD方法通过图注意力网络(GAT)[50]利用水下图像中的局部信息,使学生模型能够重建被遮挡的图像块。这种方法使模型在蒸馏过程中能更好地利用水下图像中存在的上下文相似性。
具体而言,在本文中,作者使用通过LoRA方法[51]在UIIS10K数据集上微调的SAM图像编码器作为教师模型,将小型ViT图像编码器作为学生模型。随后,作者使用学生模型在第
层输出特征
中的每个token作为图
的节点,该图节点数量不会过多,因为ViT会对输入图像进行16倍下采样。接下来,作者基于节点之间的余弦相似度连接边,以确保信息能够在相似的图像块之间传播。具体来说,对于给定的节点
,其一阶邻域
可以表示为:
其中
表示L2范数,
是一个小的值以防止除以零,
是相似度阈值。在本文中,作者使用动态阈值,将
设置为节点
与最相似的
个节点之间的余弦相似度,以防止节点连接到过多的或过少的边。
然后,作者随机遮盖了图G中大部分的节点特征,试图通过GAT重建特征
,并将其与SAM图像编码器的第
层特征
对齐,从而迫使学生网络在蒸馏过程中关注图像中的关键信息。在使用GAT重建时,作者首先通过以下公式计算被遮盖节点
与节点
之间的注意力权重:
其中
表示拼接操作,共享矩阵
是可学习的,MLP函数
用于计算节点对之间的注意力系数,
是LeakyReLU函数。重建的节点可以表示为:
最后,作者将重构的图
重形为
,并计算其与
的均方误差(MES)作为蒸馏损失。为确保学生模型和教师模型之间中间特征的逐步对齐,作者在它们的第3层、第6层、第9层和第12层应用MG-UKD。因此,最终的重建函数
可以表示为:
D. 端到端水下 Prompt 生成器
UIIS任务要求模型自动识别和分割水下图像中的每个物体。然而,现有的基础模型(如SAM [12]、SAM2 [13]或CoralSCOP [1])需要用户明确提供前景点、边界框或 Mask 作为 Prompt 来指导模型的分割。一种常见的解决方案是集成目标检测网络来识别物体的位置,并使用边界框作为这些模型的 Prompt 编码器的 Prompt 输入。然而,这种方法不仅增加了模型的复杂性,还限制了生成的 Prompt 嵌入仅包含位置信息,忽略了其他重要特征,如物体的外部外观。这种限制阻碍了模型在水下环境中的优化。相比之下,端到端水下 Prompt 生成器(EUPG)直接定位实例的位置,并将相应边界内的图像特征编码为 Prompt 嵌入
。这使得 Prompt 嵌入
能够同时包含每个实例的特征和位置信息,从而提高模型在水下任务中的性能和适应性。
EUPG的结构如图5所示。首先,EUPG通过一个简单的通道注意力函数
调整不同通道的权重,以缓解水介质引起的衰减效应,具体如以下公式所示:
其中
和
分别表示全局最大池化和全局平均池化。实际上,这个通道注意力函数可以被视为一个 Adapter [4],在图像编码器之后附加到UWSAM上,通过动态调整通道信息来增强后续生成的 Prompt 嵌入表示。
此外,为了帮助模型检测并提取不同尺寸实例的 Prompt 嵌入,作者在通道注意力调整后的特征上应用了2倍和4倍的上采样。所得特征图随后被输入到区域 Proposal 网络(RPN)[27]以实现高效的物体定位。接着,作者通过RoIAlign间隔将对应于每个实例位置的图像特征映射到
,并通过卷积、展平和MLP操作将其映射到最终特征。具体而言,上述过程可以表示为:
其中
表示位置编码,
表示RPN的位置结果。随着EUPG的引入,UwSAM的推理过程可以表示如下:
其中
表示对应实例的分类结果 因此 UwSAM 可以通过 EUPG 模块实现端到端的效率水下实例分割而无需使用 SAM 的 Prompt 编码器
E. 损失函数
作者定义作者的训练损失函数如下:
其中
是对应实例分割的损失,
是在利用MG-UKD算法进行蒸馏过程中产生的重建损失。超参数
控制这两个组件之间的平衡,在本工作中设置为
。
UWSAM的任务损失
与Mask RCNN [6]相似,由多个分量组成,即定位损失、分类损失和分割损失。因此,任务损失
可以表示为:
在EUPG中,
和
是关于分类的交叉熵(CE)损失和关于定位的Smooth L1损失的合计,而
是SAM解码器中的CE损失。
作者在UHIS10K、UIIS [2]和USIS10K [4]数据集上,使用作者提出的UwSAM方法,全面评估了当前最先进的实例分割方法,以分析其在水下目标检测和实例分割任务中的性能。对于评估指标,作者参考了MaskRCNN [6],并使用标准的bbox AP指标(
)和mask AP指标(
)来评估模型的检测和实例分割能力。此外,作者在UIIS10K数据集上与其他知识蒸馏方法进行了对比实验。遵循MGD [38]和EfficientSAM [16],作者使用Small、Medium和Large目标的mAP(
3
3
)作为评估指标。最后,作者还对其他模块进行了消融研究,以证明所提出方法的有效性。
表3 在UIIS1OK数据集上的定量比较。VIT-HUGE表示在训练期间冻结了主干网络,而LORA[52]和ADAPTER[53]分别指在训练期间应用于相应主干网络的参数高效微调方法。
A. 实现细节
作者使用PyTorch和MMDetection [62]框架实现了UWSAM以及其他对比算法。所有方法和对比算法的主干网络及超参数均与原始论文相同,仅除外作者新设计的部分。在MK-UKD中,作者采用ViT-Small作为学生网络的主干网络,并使用LoRA微调的SAM图像编码器(ViT-Huge)作为教师模型的主干网络。遵循MGD [38],作者在蒸馏过程中将特征 Mask 比例设置为0.65,并采用具有四个注意力头的双层图注意力网络(GAT)[50]来重建被 Mask 的特征。对于图构建,作者遵循WaterMask [2]并设置公式3中的
。
为了模型优化,作者在2块NVIDIA 4090 GPU上使用AdamW优化器进行训练,共训练了24个epoch,初始学习率为
,权重衰减为5e-2。在数据扩展方面,作者所有方法中都使用了随机翻转、随机缩放和随机裁剪,以确保公平比较。所有对比算法都在水下数据集上重新训练,以适应水下环境。
III-A 使用8083张图像进行训练,2010张图像进行验证和测试。
UIIS数据集。UIIS数据集[2]是首个大规模通用水下图像实例分割数据集。该数据集包含约25,000张从不同领域收集的图像,通过水下彩色图像质量评估(UCIQE)[63]和水下图像质量测量(UIQM)[64]指标筛选,选取了4,628张高质量水下图像。这些图像被精细标注为七个类别,其中3,937张用于训练,691张用于验证和测试。
USIS10K数据集。USIS10K数据集[4]是首个大规模水下显著实例分割数据集,包含7个水下类别的10,632张像素级标注图像。该数据集按照7:1.5:1.5的比例划分为训练集、验证集和测试集,分别得到7,442张训练图像、1,594张验证图像和1,597张测试图像。在本工作中,作者在USIS10K数据集上评估了yinqing-UWSAM_2505以及当前最先进的实例分割方法,以评估每种方法的泛化能力。
B. 数据集
作者在三个数据集上进行了实验:作者的UIIs10K数据集、UIIS数据集[2]和USIS10K数据集[4]。前两个数据集用于水下实例分割,而最后一个数据集用于水下显著实例分割。下面作者将详细描述每个数据集。
UIIS10K数据集。作者使用UIIS10K数据集在 underwater 环境中评估yinqing-UWSAM_2505和最先进的实例分割方法。作者遵循第 2 节中的设置。
C. 主要结果
与实例分割模型的比较。在本节中,作者评估UwSAM与最先进的实例分割方法,包括单阶段方法[56]、[58]-[60]、双阶段方法[2]、[6]、[54]、[55]、[57]、[61]以及基于 Query 的方法[30]、[32]。首先,表3展示了yinqing-UWSAM_2505与其他模型在UIIS10K数据集上的性能比较。与实例分割的通用 Baseline Mask R-CNN[6]相比,蒸馏后的UWSAM-Student模型在mAP^b和mAP^s上分别提升了3.0 AP和2.9 AP。这一改进表明MU-UKD通过知识蒸馏增强了UWSAM-Student对水下图像的特征提取能力。此外,与成熟的基于 Query 的实例分割方法YOLCAT[56]相比,UWSAM-Student在mAP、AP_50^s和AP_75^s上分别领先3.7 AP、3.4 AP和4.3 AP。这一优势源于SAM框架带来的强大实例分割能力。
表4 UIIS数据集的定量比较。VIT-HUGE*表示在训练期间冻结了主干网络,而LORA [52]和ADAPTER [3]分别指应用于相应主干网络的训练过程中的高效参数微调方法。
此外,与专为水下环境设计的两阶段方法WaterMask相比,UwSAMStudent在mAPb、
和
指标上分别领先6.3 AP、4.3 AP和6.7 AP,在mAPs、AP5o和AP5指标上分别领先1.3 AP、2.7 AP和0.4 AP。这些结果表明,与WaterMask相比,UWSAM-Student在水下实例目标检测方面取得了更优异的性能,这主要归功于EUPG模块引入的强大定位能力。最后,值得注意的是,在所有对比方法中,UWSAM-Student的参数数量是第二少的,仅为47.39 M。这使得它更适合部署在内存受限的边缘计算设备上,例如自主水下航行器和远程操控水下机器人。
作者在表4和表5中分别展示了UwSAM在UIIS数据集和USIS10K数据集上的性能表现。如表4所示,UWSAM-Student在mAPb指标上比Mask R-CNN和水Mask分别高4.0 AP和5.3 AP,在mAPs指标上分别高1.9 AP和0.5 AP。类似地,如表5所示,UWSAM-Student在水下显著实例分割任务中表现出色。具体而言,它在mAPb指标上比Mask R-CNN和水Mask分别高2.9 AP和5.6 AP,在mAPs指标上分别高1.4 AP和1.1 AP。这证明了UWSAM-Student在不同数据集和任务上的泛化能力。
与Segment Anything模型的比较。在本节中,作者将UWSAM与SAM [12]、SAM2 [13]和EfficientSAM [16]进行比较。由于SAM及其变体在推理时需要用户提供点或边界框 Prompt ,作者遵循RSPrompter [65]和USISSAM [4]中的方法,设计了两个基于SAM的流程,用于与UWSAM进行比较。具体而言,FasterR-CNN+SAM引入了一个额外的Faster R-CNN进行目标检测,其输出随后被输入到SAM中进行实例分割。同时,MaskR-CNN+SAM用SAM图像编码器替换了Mask R-CNN的主干,使得目标检测和实例分割组件能够共享一个主干进行联合优化。类似地,MaskR-CNN+SAM2和MaskR-CNN+EfficientSAM遵循与MaskR-CNN + SAM相同的原理。此外,由于USIS-SAM [4]支持端到端的实例分割,作者没有对其进行额外的修改。
表5 USIS1OK 数据集上的定量比较。VIT-HUGE* 表示在训练过程中冻结了主干网络,而 LOR! [52] 和 ADAPTER [53] 指的是在训练过程中应用于相应主干网络的参数高效微调方法。
如表3所示,与FasterR-CNN
SAM相比,UWSAM-Student在
和
指标上分别提升了10.4和3.3个百分点。这主要是因为UwSAM通过将目标检测直接聚合到模型内部,并与实例分割联合优化,从而提升了模型的性能。与专为水下环境设计的SAM变体USIS-SAM [4]相比,UWSAM-Student在
上实现了相当的性能,但其参数量仅约为USIS-SAM的7%。此外,在UIIS10K和USIS10K数据集[4]上,UWSAM-Student在
上分别比USIS-SAM高出了5.6 AP和2.2 AP。这一性能优势归因于MGUKD,它通过知识蒸馏增强了UWSAM-Student,以及EUPG,它为网络提供了强大的定位能力。与MaskR
相比,参数数量相似的UWSAM-Teacher在mAPb上高了1.7 AP,在
表6 MG-UKD与其他知识蒸馏方法在UIIS1OK数据集上的比较,其中W/O KD表示没有知识蒸馏。
表7 EUPG、W/O CA 的消融研究表明 EUPG 中的通道注意力层未被使用。
在mAPs中提升AP。这一改进主要归因于作者将SAM的 Prompt 编码器在 Pipeline 中修改为EUPG,使其能够同时执行定位和 Prompt 编码。因此,模型可以在训练过程中以端到端的方式进行优化。与 Baseline MaskRCNN相比,UWSAM-Teacher在三个数据集的mAPb上分别领先7.4、7.1和7.0 AP,在mAPs上分别领先8.8、7.5和6.4 AP。这意味着在不需要实时的一些离线水下分析任务中,UWSAM-Teacher具有巨大的潜力。
D. 消融实验
为了系统地评估UWSAM中核心组件的贡献,作者在UIIS10K数据集上进行了消融实验。
基于Mask GAT的水下知识蒸馏。作者首先将MG-UKD与几种广泛使用的知识蒸馏方法进行比较,包括OFD [66]、MGD [38]以及用于EfficientSAM的SAMI [16],以验证yinqing-UWSAM_2505的有效性。此外,遵循MGD [38]和SAMI [16],作者也使用针对小、中、大目标mAP(
)作为评估指标。为了进行公平比较,作者使用在UllS10K数据集上训练的UwSAM-Teacher作为所有蒸馏方法的教师模型。每个学生模型均使用ViT-Small作为 Backbone 网络。所有模型均使用相同超参数进行训练,共训练24个epoch,以确保一致性。
如表6所示,与成熟的基于特征的蒸馏方法OFD[66]相比,MG-UKD在mAP、
和
上分别提升了0.5、0.1和0.8 AP。值得注意的是,在更严格的
指标下取得的更大增益表明MG-UKD增强了学生模型生成更准确和详细预测的能力。这种改进可归因于MG-UKD引导学生模型恢复 Mask 特征,从而提升其表征能力。此外,与MGD[38]和SAMI[16]相比(两者均遵循 Mask 重建范式[39]、[67]),MG-UKD在
、
和
上分别比MGD提升了0.6、1.0和
,比SAMI提升了1.2、0.4和
。这些增益可归因于MG-UKD中使用的图注意力网络(GAT)[50],该网络使学生模型能够通过利用水下图像中视觉相似的 Patch 来重建 Mask Patch 。这种设计有助于学生模型在蒸馏过程中更好地利用上下文相似性,从而在水下等具有挑战性的环境中提升分割性能。
端到端水下 Prompt 生成器。为了验证EUPG在模型中的有效性,作者通过移除它来进行消融研究。所有实验均在UIIS10K数据集上进行,模型使用MG-UKD进行24个epoch的蒸馏。当移除EUPG时,模型首先使用ViTDet风格的检测器[68]进行目标检测,然后将预测的边界框传递到SAM Prompt 编码器和 Mask 解码器进行实例分割。
如表7所示,这种修改导致mAPb下降1.8 AP,mAPs下降2.1 AP。这种性能下降可以归因于模型需要首先从图像特征中预测边界框,然后将其重新编码为 Prompt 特征。特征和 Prompt 空间之间的这种重复转换引入了不必要的信噪损失,导致性能下降。作者还通过移除EUPG模块中的通道注意力(CA)层来评估其贡献。结果表明,移除CA层后,mAPb下降0.6 AP,mAPs下降0.5 AP。这表明CA层通过动态加权通道增强了模型抑制水下色彩失真的能力,从而提高了模型的特征提取能力。
参考
[1]. UWSAM: Segment Anything Model Guided Underwater Instance Segmentation and A Large-scale Benchmark Dataset
点击上方卡片,关注「AI视界引擎」公众号