向AI转型的程序员都关注公众号 机器学习AI算法工程
虽然Segment Anything Model(SAM)在医学影像分割中展现潜力,但其依赖人工标注提示的问题限制了在密集细胞核图像中的临床应用,且跨域泛化能力不足。
为此,本文提出UN-SAM框架,旨在实现全自动、跨域鲁棒的细胞核分割,以解决现有SAM方法在医学应用中的关键瓶颈。

UN-SAM包含三大核心模块:多尺度自提示生成模块(SPGen)、域自适应调优编码器(DT-Encoder)和域查询增强解码器(DQ-Decoder)。

SPGen通过多尺度特征融合和动态阈值过滤,自动生成高质量掩码提示,替代传统人工标注的点或边界框提示,显著提升临床工作流效率。
DT-Encoder在SAM的ViT编码器基础上引入域通用和域特定映射,使模型能融合跨域共有特征与特定领域知识,增强特征判别力。
DQ-Decoder利用可学习的域查询向量,在解码阶段区分不同域的细胞核类型和区域,提升跨域分割精度。这些设计使UN-SAM无需人工提示即可实现端到端的自动分割,并在多域数据中表现出优异的泛化能力。
实验结果表明,UN-SAM在多个细胞核分割数据集(如DSB、MoNuSeg、TNBC等)上显著优于传统分割方法(如U-Net、nnU-Net)和现有医学SAM变体(如SAMMI、SAMUS)。

在标准和通用评估协议下,UN-SAM的实例分割AJI和语义分割Dice系数均达到领先水平,尤其在未见域(如CryoNuSeg、CellSeg)上展现出更强的泛化能力。
消融实验验证了各模块的有效性,其中SPGen模块消除了人工标注依赖,DT-Encoder和DQ-Decoder显著提升跨域性能。
UN-SAM为医学AI领域提供了一种高效、通用的细胞核分割解决方案,有望推动自动化数字病理诊断的临床落地。
一、文献概述
本文提出了一种名为UN-SAM的领域自适应自提示分割框架,用于通用细胞核图像分割,旨在解决数字病理学中细胞核分割面临的组织类型、染色方案和成像条件多样性的挑战,同时避免传统方法对人工标注提示的依赖,提升模型在不同领域的泛化能力。

1-1:研究背景与挑战
细胞核分割的重要性与挑战
细胞核分割是数字病理学的基础,但组织异质性、染色差异和密集细胞分布导致传统方法(如U-Net、Mask R-CNN)泛化能力不足,且需复杂后处理。
基于自然图像的实例分割方法难以区分相邻细胞核。
医学SAM模型(如SAMMI、Med-SA)依赖人工标注提示(如点、边界框),不适用于含大量细胞的细胞核图像,临床应用成本高。
SAM的局限性
Segment Anything Model(SAM)在医学影像中表现出潜力,但其在细胞核分割中存在两大问题:
人工提示依赖:手动标注密集细胞的提示耗时费力。
跨域泛化不足:不同组织、染色和成像条件下的性能下降显著。
1-2:UN-SAM框架设计
UN-SAM通过三个核心模块实现全自动、跨域鲁棒的细胞核分割:
多尺度自提示生成模块(SPGen)
功能:自动生成高质量掩码提示,替代人工标注。
技术:利用特征金字塔网络(FPN)融合DT-Encoder不同层的多尺度图像嵌入,通过二值交叉熵损失优化前景区域预测,动态过滤低置信度区域,生成自提示令牌。
优势:消除人工标注需求,简化临床流程,适用于高通量分析。
领域自适应调优编码器(DT-Encoder)
功能:融合领域通用知识与特定领域知识,增强特征表达。
技术:在SAM的ViT编码器中引入“领域旁路”,通过可学习的领域通用映射( )和特定领域映射( )更新特征,冻结SAM原有参数,实现高效微调。
优势:提升模型对不同细胞核领域的适应性,减少跨域性能差异。
领域查询增强解码器(DQ-Decoder)
功能:利用可学习的领域查询( )实现跨域分割解码。
技术:通过自注意力和交叉注意力机制,结合自提示令牌和领域查询,生成特定领域的分割掩码。
优势:显式建模领域差异,提升复杂语义和实例分割的准确性。
1-3:实验与性能
数据集与评估指标
数据集:
训练/验证:DSB、MoNuSeg、TNBC、SegPC(涵盖多组织、染色和成像条件)。
泛化测试:CryoNuSeg(未知冷冻切片数据)、CellSeg(多模态细胞图像)。
指标:
语义分割:Dice系数、mIoU、F1分数、Hausdorff距离(HD)。
实例分割:聚合Jaccard指数(AJI)、检测质量(DQ)、分割质量(SQ)、全景质量(PQ)。
主要结果
跨域性能优势:
在标准和通用评估中,UN-SAM在AJI、Dice等指标上显著优于传统方法(如nnU-Net、HoVer-Net)和现有医学SAM模型(如SAMMI、Med-SA)。
例如,在DSB数据集的实例分割中,UN-SAM的AJI达78.75%,远超SAMMI(63.19%)和PromptNucSeg(76.36%)。
泛化能力:
在未知域CryoNuSeg上,UN-SAM的Dice系数为80.42%,显著优于其他模型(如Med-SA的78.46%)。
在多模态CellSeg数据集上,UN-SAM对不同成像模式(如明场、荧光)的细胞分割表现鲁棒。
消融实验:
各模块(SPGen、DT-Encoder、DQ-Decoder)均对性能有显著贡献,联合使用时Dice系数最高提升13.2%。
1-4:结论与意义
创新点:
首次实现全自动、无人工提示的细胞核分割,突破SAM的标注依赖瓶颈。
通过领域自适应设计,显著提升模型在跨组织、染色和成像条件下的泛化能力。
临床价值:
减少病理学家的标注负担,适用于大规模细胞核图像分析。
为通用医学影像分割提供新范式,有望推动AI在数字病理学中的实际应用。
1-5:代码开源
https://github.com/CUHK-AIM-Group/UN-SAM
二、整体架构
为便于理解,展示的是UN - SAM处理四种细胞核图像域的情况,且循环内和阻塞图像由来自相同域的随机小批量图像组成 2-1:DT - Encoder
Domain - Tuned Image Encoder,域调优图像编码器
输入:接收细胞核图像,先经过Patch embedding(补丁嵌入)处理 。
处理:通过Transformer Layer进行特征提取,生成图像特征表示 。图中用不同颜色标识了可训练(Trainable,橙色火焰标识 )、冻结(Frozen,蓝色雪花标识 )、循环内(In - loop,绿色锁标识 )和阻塞(Blocked,红色锁标识 )的部分。其中一些层的参数是冻结的,不参与训练,而部分是可训练的。输出的特征表示为 ,用于后续模块。
2-2:SPGen
Spatial Information Generator,空间信息生成器
输入:接收来自DT - Encoder的特征表示 。
处理:
首先进行Multi - scale Fusion(多尺度融合) ,将不同尺度的特征进行融合,以获取更全面的空间信息。
融合后的信息经过Thresh(阈值处理) ,得到自提示(Self - prompt)相关信息。这里会计算损失 ,用于优化该模块。
最终生成的自提示信息 ,一方面用于辅助分割,另一方面也输入到DQ - Decoder中。
2-3:DQ - Decoder
Domain - specific Mask Decoder,域特定掩码解码器
输入:接收来自DT - Encoder的特征 、经过处理的 ,以及域查询(Domain query) 。
处理:
内部通过Self Att.(自注意力)、Cross Att.(交叉注意力)、MLP(多层感知机)和Reverse Cross Att.(反向交叉注意力)等模块,对输入信息进行处理,捕捉不同域的特征和关联。
处理后的信息经过额外的Cross Att.和MLP进一步处理,结合自提示信息 ,最终生成分割掩码。
在这个过程中,计算分割损失 ,用于优化整个模型,以提高分割的准确性。
总体而言,UN - SAM通过这三个模块的协同工作,在无需人工标注的情况下,实现对不同域细胞核图像的高效分割。
三、项目细节
UN-SAM(通用细胞核分割任意模型)是一种域自适应自提示分割系统,专为通用细胞核图像分割任务设计。该系统在Segment Anything Model(SAM)框架基础上,引入域自适应和自提示功能的专用组件,实现了跨多种医学成像场景的高效细胞核分割。
3-1:目标与范围
UN-SAM系统旨在解决不同域(多种组织类型、染色方法或成像模态)医学图像中细胞核分割的核心挑战。
系统引入以下关键技术创新:
处理不同细胞核图像域的域自适应机制
消除手动提示需求的自提示能力
跨不同医学图像类型的通用分割性能
3-2:系统架构
UN-SAM系统由以下关键组件组成,架构图如下所示:
UN - SAM系统的架构图 UNSAM Model
DTEncoder(Domain - Tuned Image Encoder):域调优图像编码器,基于域条件处理输入图像,使模型能适应不同的细胞核图像域。
PromptEncoder:对提示进行编码,在UN - SAM中,自提示系统生成的提示会通过它编码,用于分割。
DQDecoder(Mask Decoder):域特定掩码解码器,结合图像嵌入和提示嵌入,生成分割掩码。
SPGen(Spatial Info Generator):空间信息生成器,实现自提示功能,自动生成用于分割的空间信息,无需人工标注提示。
Data Processing
JsonGen(Dataset Splitter):数据集拆分工具,负责将数据集划分为训练集、验证集和测试集 。
BinaryLoader(Data Loader):数据加载器,用于加载二进制掩码和图像数据,为模型训练和推理提供数据支持。
Training & Evaluation
train.py:主训练脚本,实现域自适应功能,使用指定的数据集对模型进行训练,优化模型参数以适应细胞核分割任务。
eval.py:评估脚本,用于评估模型在不同数据集上的分割性能,通过多种评估指标衡量模型效果。
3-3:数据流与处理
下图展示了UN-SAM系统推理过程中的数据流:
UN - SAM系统中域自适应和自提示机制相关的工作流程图 域自适应(Domain Adaptation)部分
输入:将
domain_seq
(域序列)和Input Image
(输入图像)输入到DTEncoder
(域调优图像编码器) 。domain_seq
用于指示图像所属的域相关信息,帮助模型适应不同类型的图像域。编码:
DTEncoder
对输入图像进行处理,生成image_embeddings
(图像嵌入) 。这些图像嵌入包含了图像的特征信息,且经过了域相关的调整。解码与预测:
image_embeddings
被输入到DQDecoder
(域特定掩码解码器) 。同时,自提示机制生成的prompt_embeddings
(提示嵌入)也会输入到DQDecoder
。DQDecoder
结合这两种嵌入信息,生成Predicted Segmentation Masks
(预测分割掩码) 。
自提示机制(Self - Prompting Mechanism)部分
空间信息生成:
DTEncoder
输出的image_embeddings
会输入到SPGen
(空间信息生成器) 。SPGen
利用这些图像嵌入生成spatial_maps
(空间图) 。粗分割:
spatial_maps
用于生成Coarse Segmentation
(粗分割) ,得到初步的分割结果。二值掩码与提示编码:将粗分割结果转换为
binary_mask
(二值掩码) ,然后输入到PromptEncoder
中进行编码,生成prompt_embeddings
,该提示嵌入最终会参与到DQDecoder
生成预测分割掩码的过程中。
总体而言,该流程图展示了UN - SAM系统如何通过域自适应和自提示机制,实现对输入图像的细胞核分割,无需依赖人工标注提示,并且能适应不同的图像域。
3-4:域自适应机制
UN-SAM通过域特定处理路径实现域自适应:
UN - SAM系统中域自适应处理的流程图 输入图像:流程从
Input Image
(输入图像)开始,这是待分割的医学图像等数据。域选择:图像进入
Domain Selection
(域选择)模块,通过domain_seq
(域序列)参数判断图像所属的域 。这里展示了四种可能的域,分别对应domain_seq = 1
、domain_seq = 2
、domain_seq = 3
和domain_seq = 4
。域特定处理:根据
domain_seq
的值,图像被导向不同的域特定处理路径,即Domain 1 Processing
、Domain 2 Processing
、Domain 3 Processing
和Domain 4 Processing
。每个路径会针对相应域的图像特点进行处理,例如不同组织类型、染色方法或成像模态的图像会有不同的处理方式。通用处理:经过域特定处理后,图像进入
Common Processing
(通用处理)阶段 。此阶段进行一些共享的处理步骤,可能是整合不同域处理后的特征等操作,以统一处理结果。分割输出:最后,经过通用处理的图像得到
Segmentation Output
(分割输出),即最终的细胞核分割结果。
该流程图体现了UN - SAM系统通过域自适应机制,能够根据输入图像的不同域特性进行针对性处理,同时又通过通用处理步骤整合结果,从而实现跨域的准确细胞核分割。
3-5:训练流程
UN-SAM的训练过程遵循以下步骤:
从用户启动训练到最终得到训练好的模型的整个过程 起始与参数输入
用户使用
train.py
脚本并传入参数(dataset
:数据集、sam_pretrain
:SAM预训练权重、domain_num
:域数量 )来启动训练过程。数据加载与模型初始化
train.py
调用BinaryLoader
加载训练数据。使用预训练的SAM权重初始化
UNSAM
模型,为后续训练奠定基础。
训练循环
Epoch循环:对于每个训练周期(epoch) ,进入内部的批次(batch)循环。
Batch循环:
BinaryLoader
为每个批次提供图像和掩码数据。将图像数据与
domain_seq
(域序列)一起输入UNSAM
模型进行前向传播(Forward pass) ,得到预测掩码(Predicted masks)。使用
BinaryMaskLoss
计算掩码损失(mask losses)和提示损失(hint losses) ,并将它们合并为组合损失(Combined loss) 。根据组合损失更新
UNSAM
模型的参数(Update parameters),以优化模型。
验证与模型保存
在每个epoch结束后,
BinaryLoader
加载验证数据。使用验证数据评估
UNSAM
模型的性能(Validate performance) 。根据验证性能,保存性能最佳的模型(Save best model) 。
训练完成
训练结束后,将训练好的模型返回给用户(Return trained model) ,用户可使用该模型进行细胞核分割等任务。
如需训练UN-SAM,使用以下命令:
python train.py --dataset data/$YOUR_DATASET_NAME$ --sam_pretrain pretrain/$SAM_CHECKPOINT$
训练过程使用预训练的SAM权重初始化模型,基于指定数据集实现域自适应,并针对细胞核分割任务优化模型。
3-6:技术要求
UN-SAM的依赖项如下:
Python 3.10
PyTorch ≥ 1.10.0
albumentations 1.5.2
monai 1.3.0
pytorch_lightning 1.1.0
该系统基于Segment Anything Model(SAM)框架构建,并针对医学成像应用(尤其是细胞核分割)进行了重大调整。
机器学习算法AI大数据技术
搜索公众号添加: datanlp
长按图片,识别二维码
阅读过本文的人还看了以下文章:
整理开源的中文大语言模型,以规模较小、可私有化部署、训练成本较低的模型为主
基于40万表格数据集TableBank,用MaskRCNN做表格检测
《深度学习入门:基于Python的理论与实现》高清中文PDF+源码
2019最新《PyTorch自然语言处理》英、中文版PDF+源码
《21个项目玩转深度学习:基于TensorFlow的实践详解》完整版PDF+附书代码
PyTorch深度学习快速实战入门《pytorch-handbook》
【下载】豆瓣评分8.1,《机器学习实战:基于Scikit-Learn和TensorFlow》
李沐大神开源《动手学深度学习》,加州伯克利深度学习(2019春)教材
【Keras】完整实现‘交通标志’分类、‘票据’分类两个项目,让你掌握深度学习图像分类
如何利用全新的决策树集成级联结构gcForest做特征工程并打分?
Machine Learning Yearning 中文翻译稿
斯坦福CS230官方指南:CNN、RNN及使用技巧速查(打印收藏)
中科院Kaggle全球文本匹配竞赛华人第1名团队-深度学习与特征工程
不断更新资源
深度学习、机器学习、数据分析、python
搜索公众号添加: datayx