【论文阅读】Segment Anything for Microscopy

摘要

(问题)尽管为此目的开发了许多工具,但显微图像中对象的准确分割仍然是许多研究人员的瓶颈。
(主角)在这里,我们介绍了Segment Anything for Microscopy(μSAM),这是一种用于多维显微镜数据分割和跟踪的工具。它基于Segment Anything,一种用于图像分割的视觉基础模型。
(方法)我们通过微调光学和电子显微镜的通才模型来扩展它,这些模型明显提高了广泛成像条件下的分割质量。
(扩展)我们还在napari插件中实现了交互式和自动分割,可以加快不同的分割任务,并为不同显微镜模式的显微镜注释提供统一的解决方案。
(作用)我们的工作构成了视觉基础模型在显微镜中的应用,为使用一小部分强大的深度学习模型解决该领域的图像分析任务奠定了基础。

提取:基于SAM的可交互的快速的显微图像分割模型

Intro

(背景)识别显微图像中的对象,如光学显微镜(LM)中的细胞和细胞核或电子显微镜(EM)中的细胞和细胞器是生物图像分析的关键任务之一。各种各样的模态和不同的维度(二维或三维、时间)使这些识别任务具有挑战性,迄今为止需要不同的方法。
(现有方法)最先进的方法是基于深度学习的,在过去的几年里,它们极大地改善了LM 1 -3中的细胞和细胞核分割,EM 4 -7中的细胞、神经元和细胞器分割以及LM 8,9中的细胞跟踪。这些方法中的大多数都提供了预先训练的模型,并为与其训练数据类似的新数据生成高质量的结果。
(现有方法问题)然而,由于底层深度学习方法的泛化能力有限,与原始训练数据不同的数据质量会下降,只能通过重新训练来改善。生成用于再培训的注释依赖于手动工作并且是耗时的。一些基于手动校正初始分割结果的半自动注释方法存在。如果最初的结果质量不高,这些仍然是耗时的。此外,还缺少一种统一的方法来解决LM和EM等不同模态中的各种分割任务。

(ViT)视觉基础模型最近已被引入自然图像的图像分析任务,呼应了自然语言处理的发展10 -12。这些模型基于ViT,并在非常大的数据集上进行训练。它们可用作不同分析任务的灵活主干。CLIP 10是最早成功的视觉基础模型之一,它结合了图像和语言,是许多生成式图像模型的基础14。最近,针对细分市场的基础模型已被引入11,12。其中包括Segment Anything Model 11(SAM),它是在大型标记数据集上训练的,在广泛的图像域中实现了令人印象深刻的交互式分割性能。迄今为止,这些基础模型在显微镜中的应用还很有限,但它们在该领域的潜力已经得到了确认15。

(本文方法)在这里,我们介绍了Segment Anything for Microscopy,在下文中称为μSAM,它改进并扩展了SAM的显微数据。我们的主要贡献是:

  • 一个微调SAM的训练过程,包括一个新的解码器,它提供了改进的实例分割结果。
  • 改进的LM和EM分割模型在各自领域的性能明显优于默认SAM模型。
  • 一个用于交互式和自动数据注释的工具,作为napari16插件提供。该工具可以使用默认SAM模型、我们的LM和EM模型或用户微调的自定义模型。

图1显示了μSAM的高级概述以及改进分割结果的示例。先前的工作已经研究了SAM的生物医学应用,例如,在医学成像17,组织病理学18和神经成像19。然而,这些研究仅限于默认SAM,并没有对各自的领域进行再培训,根据我们的研究结果,这是至关重要的。已经针对医学图像数据中的窄交互式分割任务研究了用于其他域的重新训练SAM 20。使用SAM作为自动分割的基础已经被研究用于组织病理学21,并且使用它用于细胞分割已经基于先验对象检测被研究22。这些先前的工作都没有像我们的贡献那样在单个模型中结合重新训练完整的交互式分割功能和改进的自动分割。

(效果对比)与现有的分割和跟踪工具相比,μSAM更通用,因为它的预训练模型涵盖了LM和EM,涵盖了广泛的分割任务。它支持二维(2D)和体积分割以及在同一工具中进行跟踪。它结合了使用相同的基础模型的交互式和自动分割。因此,模型的两个方面都在微调过程中得到了改进,这可以大大加快数据注释的速度。相比之下,CellPose 2(参考文献23)的在环训练模式(开创了集成数据注释和训练)依赖于手动像素级校正。总之,我们的工具的主要区别特征是它适用于不同模态和维度的各种分割任务,以及由于其交互式分割功能而具有的快速注释速度。我们在三项用户研究中展示了这些方面,我们发现CellPose23在细胞分割方面具有竞争力的性能,与ilastik carving24相比,体积分割的性能明显提高,与TrackMate9相比,跟踪性能明显提高。总的来说,我们的贡献显示了视觉基础模型在生物成像中统一图像分析解决方案的前景。我们的工具可在https://www.example.com上获得。github.com/computational-cell-analytics/micro-sam/

在这里插入图片描述

Results

我们将默认SAM与针对不同显微镜分割任务进行微调的模型进行比较。
(数据集)首先,我们在LIVECell25数据集上研究了交互式和自动分割。然后,我们训练和评估通才模型,包括在多个数据集上进行训练,用于LM中的细胞和细胞核分割以及EM中的细胞核和细胞核分割。
(三个模型)在下文中,我们将Kirilov等人提供的原始模型称为**“默认”模型**,将我们在单个数据集上微调的模型称为**“专家”模型**,将我们在多个数据集上微调的模型称为**“通才”模型**。
(存在的问题和通才模型)请注意,考虑到当前的SAM架构,训练一个在不同显微镜模式下持续改进的单一模型是不可行的(有关详细信息,请参见“EM”部分和“讨论”)。因此,我们为LM和EM训练单独的通才模型。我们进一步研究微调SAM在资源受限的设置。
(交互工具)然后,我们介绍我们的用户友好的工具,实现为napari16插件,用于(体积)分割和跟踪的交互式和自动数据注释。我们将其与LM中的细胞分割、EM中的细胞核分割和LM中的细胞核跟踪的三个用户研究中的既定工具进行比较。

Fine-tuning SAM improves cell segmentation

(SAM介绍)在参考文献11中,SAM被引入作为交互式分割的模型:它基于点、框或掩模注释预测对象掩模。点注释可以是正的(对象的一部分)或负的(不是对象的一部分)。该模型是在一个非常大的带有对象注释的自然图像数据集上训练的。作者还介绍了一种自动实例分割方法,称为自动掩模生成(AMG),该方法基于用点网格覆盖图像,使用所有点作为SAM的点注释,并过滤掉不太可能或重叠的掩模。他们评估了各种任务的交互式和自动分割,包括LM机器人26。有关SAM功能的概述,请参阅方法。
(SAM用于显微图像的问题)原始的显微镜实验和我们对默认SAM的评估显示出对显微镜的显著推广,尽管原始训练集主要包含自然图像。然而,我们注意到显微镜模型的几个缺点。例如,SAM将细胞簇分割为单个对象,如图1b所示。
(改进方法)为了改进SAM应用于我们的领域,我们实施了一个迭代训练方案,以实现对新数据集的微调。这种方法重新实现了原始的训练方法,到目前为止还没有开源。此外,我们在模型中添加了一个新的解码器,预测前景以及到对象中心和边界的距离,然后通过后处理获得自动实例分割。我们将这种方法称为AIS。附加的解码器可以与SAM的其余部分一起训练。AIS和培训方法在方法(“AIS”和“培训”)中有更详细的解释。

(细胞分割 微调介绍)我们在LIVECell 25上研究了我们的微调方法,LIVECell 25是最大的公开可用的细胞分割数据集之一。
(数据集表现)图2a显示了默认和LIVECell专家模型的平均分割准确度27(越高越好;详情请参见方法),使用单独的测试集进行评估。在这里,我们通过基于分割地面实况模拟用户注释来评估交互式分割。我们从地面真值中导出一个框注释(红色条)或一个正的点注释(绿色条),对应于图中的迭代0。然后,我们从预测中不正确的区域采样一个正值和一个负值点,正值点是预测缺失的区域,负值点是预测不应该缺失的区域,然后使用附加注释重新运行模型。该过程重复七次(迭代1-7),并且在每次迭代中,新采样的点被用作附加的点输入。
我们还比较了通过AMG和AIS(仅在微调后可用)的自动分割,并提供了在LIVECell上训练的CellPose模型的结果以供参考。我们使用基于大型视觉Transformer(ViT-L)的SAM,并在LIVECell的训练分割上训练它250,000次迭代。我们发现,ViT-L在运行时间和质量之间提供了最佳平衡;不同模型大小的运行时间比较见图5a,分割结果评估见扩展数据图1。结果显示,由于在所有设置上进行了微调,因此性能得到了明显改善。与专业模型的交互式分割明显优于任何自动分割结果,而它只有在对默认模型进行多次校正迭代后才能达到CellPose的性能。

当提供更多标注时,专家模型也实现了一致的改进,而默认模型不是这种情况。这部分是因为我们不使用掩码预测作为附加模型输入;详情见补充图1。AMG的实例分割性能大幅提升,而AIS的分割性能(仅在微调后才可用)与CellPose不相上下。

(微调策略比较)我们在图2b中研究了不同的微调策略,其中我们只微调SAM架构的一部分,冻结所有其他权重。这里,我们进行与图2a相同的评估实验。为了更简洁地展示,我们仅报告单点提示(“点”)、单框提示(“框”)以及从点提示开始时的最后一次迭代的交互式分割结果(“IP”,对应于图2a中迭代7处的绿色条)以及从框提示符开始时的最后一次迭代(“IB”,对应于图2a中迭代7处的红色条)。结果表明,对图像编码器进行微调对算法的影响最大,而对完整模型进行微调的整体性能最好。在图1c中,我们使用LIVECell出版物中定义的数据拆分,仅使用可用训练数据的一个子集对模型进行了微调。实验结果表明,训练数据分数分别为2%、4%和5%时,训练效果最好。总体而言,LIVECell的结果提供了以下结论

  1. 微调SAM可明显提高给定数据集的分割质量。
  2. 微调模型的所有部分会产生最佳结果。因此,我们在所有进一步的实验中训练完整的模型。
  3. 对于给定数据集的大部分改进可以用训练集的相当小的一部分来实现。我们在图5b中对此进行了更详细的研究。

An LM generalist model improves across diverse conditions(LM通才模型)

我们的下一个目标是为LM训练一个通用模型,该模型可以提高该模态的分割性能,从而可以替代默认SAM。
虽然之前的实验已经表明,对给定图像设置的数据进行微调可以提高性能,但我们还没有证明它可以提高泛化能力。
(训练方法)为了训练通才模型,我们根据已发表的数据集(包括LIVECell25、DeepBacs28、TissueNet2、NeurIPS CellSeg29、PlantSeg(Root)30和Nucleus DSB26),使用该数据集的一个版本(不包括StarDist3提供的组织病理学图像)和来自细胞跟踪工具31的八个数据集,组装了一个大型且多样化的训练集。我们还在其中五个单独的数据集上训练专家模型。图3a比较了默认、专家和通才模型的细分性能。在所有情况下,评估都是在不用于训练的测试分割上完成的。与默认模型相比,我们看到专家和通才模型都有明显的改进。通才模型总体上表现类似于或优于专家,除了LIVECell上的AIS。我们包括CellPose的自动分割结果作为参考,使用LiveCell和TissueNet的专业模型以及“cyto2”模型。通过AIS自动分割的性能与CellPose相当或更好,但TissueNet除外。我们认为,这种差异部分是由于TissueNet包含两个通道,这两个通道不能很好地映射到SAM(方法)的RGB输入。请注意,DeepBacs,PlantSeg(Root)和NeurIPS CellSeg的比较严重偏向于我们,因为我们的模型是在各自数据集的训练分割上训练的,与CellPose cyto2模型不同。

(泛化能力)为了研究通才模型是否可以提高对其他显微镜设置的泛化能力,我们将其应用于训练集中没有直接表示的数据集。我们选择包含免疫荧光数据的数据集COVID IF 32、包含用共聚焦荧光显微镜成像的植物细胞的PlantSeg(胚珠)30、包含组织病理学图像的Lizard 33和包含用共聚焦荧光显微镜成像的哺乳动物细胞的Mouse Embryo 34。默认SAM和我们的通才模型以及CellPose中的模型cyto 2的结果也显示在图3a中。对于所有数据集,我们的通才对默认模型的改进是显而易见的。AIS在几乎所有情况下都优于AMG,并且CellPose和AIS显示出总体上相当的分割精度。对于包含与训练数据和小鼠胚胎不同的模态的Lembryo,由于细胞尺寸小,这是一个特别困难的问题,没有一种自动分割方法效果良好,但交互式分割产生了良好的结果,并由通才进行了改进。图3b显示了默认模型和通才模型的分割比较示例。我们报告了ViT-L的结果,为我们的模型训练了250,000次迭代。扩展数据图2显示了不同大小的SAM模型的结果,包括其他数据集的结果。图3和4示出了定性示例,补充图2b示出了自动分割结果的示例。LM数据集的概述见补充表1。

总体而言,这些实验表明,给定领域的通才模型明显提高了分割质量。我们为LM提供了这样一个通才模型。它支持交互式和自动分割,实现了与CellPose相当的自动分割质量,CellPose是自动细胞分割的最新技术。请注意,我们并不声称我们的模型在自动分割方面优于CellPose,但它为大多数实际设置提供了类似的质量,同时还支持交互式分割。

Improved mitochondria segmentation in EM (EM通才模型)

我们进一步研究训练EM的通才模型。与LM相比,这更具挑战性,因为在EM中,膜结合结构被非特异性地标记,而不是具有针对细胞组分的特异性染色。因此,EM中的分割任务更加多样化,并且结构可以具有分层组成,例如,细胞室内的细胞器。这使得训练用于通用EM分割的模型更具挑战性。因此,我们专注于训练一个用于分割线粒体和细胞核的模型,其中存在大型公共数据集。我们利用MitoLab5和MitoEM35数据集用于线粒体,PlatyEM4用于细胞核。我们在下文中将此模型称为EM generalist,但要明确的是,它仅可靠地改善了线粒体、细胞核和其他圆形细胞器的EM分割。由于不能提供一个统一的模型EM的限制,我们也避免探索一个统一的通才模型EM和LM。

(线粒体数据集)我们在训练数据集的测试拆分和其他测试数据集上比较了默认的SAM和我们的EM通才模型:在FIBSEM中包含线粒体的Lucchi36,在FIBSEM中包含线粒体的两个MitoLab5测试数据集,在FIBSEM中包含线粒体的UroCell37,以及在连续切片TEM中包含线粒体的VNC。我们还包括NucMM(Mouse)38,它包含在高能X射线中成像的核,这是一种与EM相似的成像方式。定量和定性结果见图4。我们看到,由于对所有数据集进行了微调,交互式分割得到了明显的改进。对于线粒体自动分割,我们也与MitoNet5进行了比较,发现其性能总体上与AIS和AMG相当,结果因数据集而异。与AMG相比,AIS的优势不像LM那么明显。这可能是因为AMG对分离良好的物体,如EM中的线粒体,比密集堆积的物体,如LM中的细胞,工作得更好。在实践中,由于其较低的运行时间,AIS在大多数情况下是可取的(图5A)。请注意,我们并不声称我们的方法在自动线粒体分割方面优于MitoNet,而是说它提供了类似的质量,同时也实现了交互式分割。扩展数据图5显示了附加数据集和不同模型大小的结果。扩展数据图6和图7显示了其他定性结果,而补充表2列出了EM数据集的概述。

(其他细胞器数据集)我们还在EM中进行其他细胞器和结构的实验。我们用我们的模型分割纤毛和微绒毛(参见扩展数据图中海绵EM和扁虫(纤毛)的结果)。5-7)并发现我们的EM通才模型总体上比默认SAM执行得更好。我们还研究了内质网(ER)和神经突的分割(扩展数据图8)。我们发现我们的EM通才在这些情况下仅提供了边际效益或有害,这是由于ER与线粒体/细胞核相比的不同形态学以及该模型相对于周围细胞隔室更喜欢分割细胞器的事实。我们针对这两种情况训练专家模型,这明显提高了给定分割任务的性能。

总之,我们发现在EM中训练用于改进的细胞器分割的模型是可行的,并且我们提供了用于线粒体和细胞核分割的通用模型,这也可以改进用于具有相似形态的其他细胞器的结果。给定合适的数据集,训练甚至更一般的EM模型应该是可能的,但是为了训练改进细胞区室和细胞器的分割的真正的通才模型,需要语义上知晓的模型和训练过程。但是,我们的微调方法可用于为给定EM分割任务训练专家模型,我们的注释工具(见下文)可用于快速数据注释,以提供所需的训练数据,这使得我们的贡献对于EM分割任务也很有价值,因为我们的EM通才模型无法提供优势。

Resource-constrained settings for inference and fine-tuning(资源限制)

我们的主要目标之一是建立一个用户友好的工具,用于交互式和自动显微镜分割。
作为一个准备,我们调查如何SAM可以用于资源受限的设置,例如,在用户的笔记本电脑或一个普通的工作站,推理和微调。
(时间比较)首先,我们比较了所有相关操作的推理时间:计算图像嵌入,使用框或点注释对一个对象进行推理,以及通过AMG和AIS对CPU和GPU进行自动分割(图5a)。对于Point,Box,AMG和AIS,我们测量不包括嵌入计算的运行时间。GPU上的运行时间要小得多,但考虑到预先计算的嵌入,在CPU上使用点或框进行交互式分割是可行的,每个对象大约需要30 ms。我们还看到AIS与AMG相比有很大的加速。GPU的主要优点是大大减少了嵌入计算的运行时间和更快的自动分割,特别是对于多维数据(见下文)。
我们还比较了不同大小的图像编码器的运行时,包括ViT Tiny(ViT-T)39。考虑到运行时和分段准确性之间的权衡(参见扩展数据图)。1、2和5,以广泛比较不同模型大小的分割质量),我们建议使用ViT Base(ViT-B)或ViT-L模型。使用ViT Huge(ViT-H)通常不会产生更好的结果,但会产生更高的计算成本。如果运行时间是一个问题,ViT-B可以使用,只有一个小的惩罚分割质量。ViT-T的速度要快得多,对于简单的分割任务来说效果很好,但对于其他任务来说质量严重下降。为了与现有工具进行比较,我们还使用相同的硬件测量了CellPose的运行时间。使用GPU分割图像大约需要0.3秒,使用CPU需要1.5秒。与此相比,嵌入计算和自动分割的相关措施AIS的运行时间为0.2-1.2 s,

我们进一步研究模型微调资源受限的设置。虽然我们的LM或EM通才模型在许多设置中提高了质量,但它们可能不足以满足用户的需求,或者可能与其数据的模态不匹配。为了能够进一步改进特定任务,我们研究了对COVID IF数据(也用于图3)的微调。我们研究了它在少量注释图像和CPU微调(图5b)以及其他硬件配置(补充图3)中的行为,从默认或LM通才模型开始。为了能够在有限的资源下进行训练,我们使用早期停止并找到最佳的超参数,以便能够针对给定的硬件配置进行训练(扩展数据图10 b,c)。我们还研究了使用LoRA 40的参数有效训练,它有望实现更快的训练(参见前面提到的图中的虚线结果线和补充图4c,以广泛评估LoRA的训练)。我们发现,使用CPU仅对少数图像进行训练是可行的,并且它可以针对给定任务明显改善模型。使用LoRA进行训练在大多数情况下会导致更长的训练时间,因为模型需要更多的迭代来收敛。在此设置下,CPU上的训练花费了5.3小时,而GPU上的训练花费了30分钟。不同硬件配置的训练时间概述见扩展数据图10d。

总体而言,我们发现在资源受限的环境中应用和培训SAM是可行的。然而,与基于卷积神经网络的架构相比,计算图像嵌入和训练的运行时间更大,特别是在使用CPU时。我们还发现,对一些带注释的图像进行微调,可以使用我们的注释工具快速生成(下一节),明显改善了结果。从我们的模型开始微调可以提供明显的好处。在我们的模型比默认SAM差的情况下,例如,对于EM中的神经突分割,从默认模型开始可能会更好,因此用户应该选择最适合其任务的模型作为起点。

μSAM enables fast data annotation for microscopy(快速数据标注)

我们提供了一个用于交互式和自动数据注释的工具,利用前面部分中描述的模型和知识。为了让生物学家可以轻松使用该工具,我们将其实现为napari 16插件。Napari是一个基于Python的多维图像数据查看器,在显微图像分析中很受欢迎。我们在我们的工具中提供了五种不同的功能:(i)用于2D图像分割,(ii)用于体积分割,(iii)用于时间序列数据中的跟踪,(iv)用于多个图像的高通量分割和(v)用于微调。它们被实现为单独的插件小部件。注释小部件(i-iv)支持基于用户提供的点或框注释的交互式分割以及基于AIS或AMG的自动分割(跟踪小部件(iii)除外)。为了实现交互式使用,我们实现了图像嵌入的预计算和缓存、分片交互式自动分割以及在参数更改的情况下自动分割的高效重新计算。我们还支持交互式分割的体积数据和交互式跟踪的时间序列投影掩模相邻的切片或帧和重新运行SAM与派生的注释。对于体积数据,我们通过在每个切片上运行AIS或AMG并在后处理步骤中合并切片之间的结果来实现自动分割(扩展数据图9)。微调小部件(v)允许用户选择最适合其硬件的模型和训练参数,然后根据自己的数据对模型进行微调。我们还提供底层功能作为Python库,以便具有计算知识的用户可以实现培训脚本,并且开发人员可以在我们对原始SAM功能的扩展基础上进行构建。详见《办法》。补充视频解释了工具的使用,并在https://computational-cell-analytics.github.io/micro-sam/micro_sam.html.

我们研究了我们的工具三个代表性的注释任务:在明亮显微图像下的类器官分割,细胞核分割在EM和细胞核跟踪荧光显微镜,并比较它们建立的软件各自的注释任务。有关用户研究实验设置的更多详细信息,请参见方法和补充信息。

User study 1: Brightfield organoid segmentation

对于二维标注,我们研究了明场图像中的类器官分割。生长类器官是研究组织的常用实验技术,例如,在癌症研究中。器官分割使得能够研究生长和形态。在这里,我们使用一个内部数据集来比较不同的注释方法,将我们的工具与CellPose和手动注释进行比较。图6a中总结了研究结果。在我们的工具中,我们使用默认SAM(“μ SAM(默认)”)、LM generalist(“μ SAM(LM generalist)”)和根据用户注释进行微调的模型(“μ SAM(微调)”)进行比较。对于所有这些模型,我们首先运行自动分割,然后使用交互式分割进行校正。我们使用ViT-B作为所有型号的图像编码器。对于CellPose,我们使用cyto 2模型(“CellPose(默认)”)、从cyto 2开始的环内训练(“CellPose(HIL)”)和在环内训练后获得的模型的注释(“CellPose(微调)”)。在这里,我们也首先运行自动分割,然后使用手动注释进行校正。使用CellPose GUI执行这些实验。对于每种方法,我们报告了每个对象的平均标注时间,在微调后,这提供了类似的标注速度和质量,但推广到类似的数据较差。我们还想强调的是,结果是数据相关的,根据初始模型的性能,其他数据集的结果也会有所不同。

User study 2: volume EM nucleus segmentation

对于三维(3D)注释工具,我们使用果蝇幼虫大脑的内部数据集研究了体积EM中的核分割,为此我们还对几个小块进行了地面实况注释。在体积EM中分割细胞核或其他大细胞器是分析细胞形态和基于表型标准区分细胞类型的重要任务4。在这里,我们比较了交互式细胞核分割与μSAM和ilastik雕刻24。雕刻使用种子图形分水岭从用户注释中分割3D对象。这种方法不是基于深度学习,但仍然是交互式3D分割最常用的方法之一,例如,Gallusser等人41使用它来生成训练注释。在μSAM中,我们首先使用默认的ViT-B模型对数据进行注释,该模型的效果略好于EM通才模型,这可能是由于我们对训练数据的分辨率存在差异。在这种情况下,我们没有使用自动分割,因为它没有产生好的结果。我们还用地面实况数据在另一个小块上微调模型。对于该模型,自动3D分割(基于AIS)产生了良好的结果。图6 b示出了每个对象的注释时间和注释块的图示。当使用默认模型时,使用μSAM的注释比ilastik更快,当使用微调模型时,我们可以纠正自动分割结果,而不是交互地分割每个对象。

User study 3: fluorescence microscopy nuclei tracking(荧光显微镜核跟踪)

我们研究了Schwartz等人42在荧光显微镜下成像的细胞核数据集上的跟踪注释工具,使用每三帧一次,使任务更具挑战性。我们将通过μSAM进行的注释与最新版本的TrackMate 9进行了比较,后者集成了对基于深度学习的分割工具的支持,包括StarDist 3。图6c显示了四种不同方法的结果:使用我们的工具进行交互式分割,使用默认SAM,LM通才模型和针对此数据进行微调的模型,以及使用StarDist的TrackMate。我们报告的注释时间和质量的注释相比,地面真相。请注意,我们的工具和TrackMate在跟踪方面的工作方式完全不同:在我们的工具中,必须以交互方式跟踪每个谱系,而TrackMate则根据StarDist的分割自动跟踪细胞核,然后进行手动校正。在这里,我们看到了LM通才模型相对于默认SAM的明显优势;它在交互式注释期间更好地跟踪细胞核。在单独的时间序列上对该模型进行微调不会进一步加快跟踪速度。与TrackMate相比,我们的方法有点慢,这是因为我们目前不会自动跟踪对象,但会产生更高质量的注释。我们的目标是实现自动跟踪,可以用作校正的起点,基于AIS的初始逐帧分割,并期望从这个扩展中获得重大的加速。

Discussion

我们已经介绍了一种方法来微调SAM显微镜数据,用它来提供LM和EM的通才模型,并广泛比较这些默认SAM和自动分割的参考方法。我们还实现了一个napari插件,用于交互式和自动分割。我们的定量实验和用户研究表明,我们的贡献可以加快数据注释和自动分割的各种应用程序。我们的贡献也标志着视觉基础模型在显微镜中的应用。我们希望未来的工作建立在它的基础上,并扩展这些模型的应用,以进一步提高对象识别任务,并解决其他图像分析问题。

我们将我们的方法与现有的分割和跟踪工具进行比较,并显示出竞争力或改进的性能。然而,我们希望通过将我们的方法的一部分与其他工具集成,可以进一步提高可用性和性能。例如,我们的模型和交互式分割功能可以与CellPose、MitoNet或其他自动实例分割方法集成,使用户能够进行微调,将更快的数据注释与更有效的架构相结合,以处理大型数据集。为了实现这种集成,我们将我们的注释工具开发为napari插件,以便它们可以与其他基于napari的软件结合使用,在BioImage.IO43上发布了我们的模型,以标准格式提供它们,并提供了一个文档齐全的Python库。我们的模型已经可以在Deep MIB 44和QuPath 45,46中使用,它们为SAM提供了初步支持。与其他支持交互式注释的工具(如ilastik24或TrackMate9)集成也是可取的。

计划)我们还计划在多个维度上改进和扩展μSAM。在不久的将来,我们计划为生物医学应用训练更多的模型,特别是利用Open Organelle7提供的数据和其他模式(如组织病理学数据)的细胞器分割的通用EM模型。我们还希望实现自动跟踪,以加快使用跟踪工具进行注释的速度。为了实现更有效的微调,我们计划将对参数高效训练方法的研究扩展到比LoRA40更新的方法,这可能会在我们的设置中提供更快的训练时间。此外,更高效的架构47可以取代基于变换器的编码器,以降低推理和训练的计算成本。为了实现显微镜实例分割的通用模型,我们计划还研究如何使SAM(或类似模型)具有语义感知,以实现EM中的模糊分割情况,如何将其扩展到完整的3D分割以及如何训练多个域(LM,EM)的统一模型。

问题)虽然我们的贡献为交互式和自动显微镜分割提供了多功能和强大的功能,但与现有方法相比,它具有一些局限性,主要是由于视觉变换器的计算占用空间较大。虽然由于SAM的模块化设计(实现图像嵌入的预计算),交互式数据注释是可能的,但与基于CNN的方法(如CellPose或MitoNet)相比,大型数据集的自动处理效率不高。此外,在新数据上微调SAM模型需要更长的时间,特别是在CPU上,因此我们不像CellPose那样提供“人在回路中”的微调,在CellPose中,模型在每个注释的图像之后更新,而是让用户通过单独的用户界面或脚本进行微调。计算成本也使我们无法构建一种在正交切片上操作的3D分割方法,正如CellPose和MitoNet所做的那样;我们逐切片处理体积或时间序列数据,并使用后处理来避免产生的伪影。由于2D推断而导致的一些伪影仍然可能发生。

免责声明:没有训练两个sota工具)我们与CellPose和MitoNet的比较旨在为用户提供自动分割工具的参考。虽然我们已经尽了最大努力与这些方法进行公平的比较,但我们没有在模型的训练数据上重新训练它们(这对于SAM最初训练的大型数据集来说是非常具有挑战性的)。我们并不声称与它们相比具有上级性能;相反,我们为大多数实际用途提供了类似的自动分割质量,并增加了交互式分割的好处,并支持更多的数据模式。同样,我们进行的用户研究有很多自由度,因此根据用户体验和用例,关于工具适用性的结论会有所不同。尽管如此,我们认为这些研究为我们的工具在实践中的应用提供了重要的背景,我们已经设计了它们,以提供尽可能公平的比较。此外,我们还没有提供一个单一的模型,同样适用于多个显微镜领域,而是提供了三套模型(LM通才,EM通才线粒体和细胞核,默认SAM)具有不同的优势。我们在文档中添加了一个部分,以指导用户为其应用程序选择正确的模型(https://computational-cell-analytics.
github.io/micro-sam/micro_sam.html#choosing-a-model
).我们相信,尽管有这些限制,μSAM提供了目前可用的(交互式)显微镜分割的最通用的解决方案,我们乐观地认为,本文概述的发展将最终解决其局限性。

任何方法、其他参考文献、Nature Portfolio报告摘要、源数据、扩展数据、补充信息、致谢、同行评审信息;作者贡献和竞争利益的详细信息;以及数据和代码可用性声明,请访问at https://doi.org/10.1038/s41592-024-02580-4

Methods

SAM

SAM是交互式细分的视觉基础模型。Kirilov等人11介绍了该方法。在此,我们简要总结一下它的主要功能。它通过预测用于描述输入图像中给定对象的注释的对象掩模来解决交互式分割任务。注释可以是边界框、点(正和/或负)或低分辨率遮罩。在SAM出版物中,这些注释称为“提示”;我们交替使用术语“提示”和“注释”。Kirilov等人也描述了基于文本注释的分割,但模型的已发布版本不包括此功能。对于新的图像,该模型利用图像编码器来预测嵌入,该嵌入对应于输入的缩减表示中的每像素的矢量。图像编码器是一种视觉转换器13,SAM有三种不同尺寸的编码器,采用ViT-B、ViT-L或ViT-H架构(按型号尺寸递增排序)。我们还包括一个使用较小的ViT-T的版本,它是由MobileSAM 39推出的。图像编码器包含模型的大部分参数。每个图像只需应用一次,这样,如果交互式分割中的注释发生变化,就可以快速重新计算对象遮罩。该模型的其他部分是对用户注释进行编码的提示编码器和基于图像嵌入和编码的注释预测对象掩码和IOU分数的掩码解码器。IOU分数对应于对掩模质量的估计。为了处理单点注释的模糊性,其可以同时涉及对象或其一部分,SAM针对这种情况预测了三种不同的掩码。SAM架构概述见扩展数据图1a。

该模型是在自然图像的大的标记数据集上训练的,该数据集是由注释者迭代构造的,该注释者校正在该数据集的先前版本上训练的SAM的输出。该模型在多种分割任务上的测试结果表明,该模型对不同领域的图像具有良好的泛化性能.作者还实现了一种用于AIS的方法,称为AMG。它使用点网格覆盖输入图像,并预测所有点的遮罩。对预测的掩码进行后处理,以仅保留高质量的预测。这涉及滤除具有低IOU预测的掩码和具有低稳定性分数的掩码,该稳定性分数是基于在不同logit值处阈值化时掩码的变化来计算的。最后,应用非最大值抑制来移除重叠预测。

SAM是在RGB图像上训练的,因此图像编码器需要三个通道的图像数据作为输入。为了处理显微镜图像,其中大多数具有单个通道,我们将该通道复制三次。我们发现这种方法效果很好,并假设SAM也使用相同的方法在灰度图像上训练。将该模型应用于具有不同数量通道的数据,例如,对于细胞核和细胞质染色(例如在TissueNet中),使用两个通道,更具挑战性。我们尝试了两种方法:(i)附加一个空声道,以及(ii)对两个声道求平均以获得一个声道,然后将该声道复制三次。这两种方法都有缺点:在第一种情况下,与通过添加空通道的训练相比,图像统计被改变,而在第二种情况下,通过求平均而丢失了信息。我们发现第二种方法效果更好,并将其应用于相关实验中。请注意,与单独使用两个通道相比,这种方法是有害的,并且在将当前SAM架构应用于多通道图像时构成了限制。在将图像传递到图像编码器之前,将图像大小调整为1,024 × 1,024像素。

AIS

我们扩展了原始的SAM结构,增加了一个用于预测AIS的解码器。该解码器是基于UNETR 48。它由两个卷积层的四个块组成,每个卷积层后面跟随用于上采样的转置卷积。每个块接收图像编码器输出作为附加输入。解码器的输出具有与输入图像相同的空间维度。它预测了三个输出通道:到对象中心的距离、到对象边界的距离和前景概率。每个对象的距离都是标准化的;关于用于训练的目标的描述,请参见补充图2a。我们使用scikit-image 49中的实现,使用种子分水岭,基于它们计算实例分割。两个距离通道都用于通过查找中心距离低于阈值参数且边界距离高于阈值参数的连通区域来导出种子。除了这些种子之外,分水岭还使用距离预测作为高度图,并使用阈值前景预测作为遮罩。我们已经选择了这种方法来分割复杂的对象形态,其过程相当简单:使用边界距离可以防止合并窄的相邻对象,如果仅使用中心距离,则会错误地连接这些对象。相反地,使用中心距离防止了在阈值边界距离预测中错误地分割具有多个连接区域的非凸对象。��于此,我们称之为AIS。该分割过程受使用距离预测进行实例分割的其他方法(例如,StarDist 3或CellPose 1)的启发,但它使用更简单的后处理逻辑。

我们通过将我们的方法与另外两种分割方法进行比较来验证它:预测边界和前景,然后是分水岭,以及预测相似度,然后是Mutex Watershed 50。我们基于LIVECell上的SAM ViT-B编码器为所有三种方法训练了UNETR模型,使用10,000次训练迭代,否则使用下一节中描述的相同超参数。我们发现,基于距离的方法(平均分割精度为0.39)比预测亲和力(0.36)和边界(0.31)更好。我们还进一步比较了使用不同网络架构时我们的分段方法的工作原理。为此,我们将UNETR架构与UNet 51和基于SAM的更简单架构进行了比较,该架构重用SAM图像编码器和掩模解码器来预测前景和距离通道以进行实例分割。结果如补充图4a、b所示。总之,我们看到,基于SAM的体系结构为小训练数据集提供了很大的优势,只要它们的权重用预训练模型初始化,并且具有卷积解码器的UNETR体系结构在此任务中比使用SAM掩码解码器更有优势。

请注意,我们的分割方法也与CellVIT21有一些相似之处,后者在组织病理学中使用SAM编码器进行AIS。但是,CellVIT不保留SAM的交互式分段功能。我们还评估了它的实例分割方法,该方法基于预测距离梯度,但发现它不适合触摸对象。

Training

为了微调SAM模型,我们按照Kirilov等人的描述实现并提供了一个迭代训练方案。请注意,Kirilov等人的训练算法到目前为止还没有发布。微调SAM的其他工具(例如MedSAM20)依赖于一种更简单的训练启发式方法,该方法仅针对特定类型的提示(例如框提示)微调SAM。我们已经发现,这样的方法提高了给定提示类型的分割质量,但是它们妨碍了其它提示的分割质量(也参见下文)。为了提供用于交互式分割的模型,因此遵循与用于训练初始SAM的训练过程类似的训练过程是至关重要的。

训练算法需要图像数据和感兴趣对象的相应地面实况分割。在训练过程中,我们在所谓的epoch中对完整的训练集进行了多次迭代。在一次迭代中,我们对一个小批量进行采样,对应于多个图像和相应的地面实况,应用图像编码器,从地面实况中导出提示,然后将其传递给提示编码器,并使用掩码解码器预测对象。然后我们计算预测值和真实值之间的损失,并通过反向传播和梯度下降更新网络权重。与用于实例分割的常规训练方法相比,单次迭代更复杂,因为它由多个子迭代组成以模拟交互式分割。更详细地说,训练迭代遵循以下步骤:

  1. 从训练集中对包含输入图像和地面实况的小型批处理进行采样。

  2. 从地面实况中采样固定数量的物件。使用给定图像中的所有对象进行训练将需要太多的内存。

  3. 用编码器预测采样图像的嵌入。请注意,编码器仅取决于图像数据,而不取决于提示符。

  4. 以批处理方式对所有采样对象执行以下步骤:
    a.从对象中随机采样一个点(用作正输入点),或使用对象的边界框作为提示。
    B.使用SAM预测指定输入的掩码和预期IOU值。如果模型以单点注记显示,则预测三个输出掩码;否则,预测单个输出掩码。有关此方法的动机,请参阅上一节。
    C.计算预测对象和地面真实对象之间的损失,以及估计IOU得分和真实IOU得分之间的损失。如果预测了三个对象,则仅考虑具有最高IOU预测的对象的丢失。
    D.示例两个新点:一个是正的,其中模型预测了背景,但是应该存在前景(根据地面实况),一个是负的,用于相反的情况。如果因为没有缺少前景预测的区域而无法对这些点进行采样,或者反之亦然,则我们从对象内/外随机采样正/负点。
    e.向模型显示前面步骤中的组合注释,即到目前为止采样的所有点和框注释(如果在第一步中使用)。我们还使用上一步中的掩码预测作为概率为50%的附加提示;有关此步骤的更多详细信息,请参见下一段。
    F.计算当前预测的掩码和IOU丢失。

  5. 步骤4d-f重复固定次数(我们使用8个子迭代),累积所有损失;基于所有子迭代上的平均损失执行反向传播和梯度下降,并更新模型的所有部分(图像编码器、提示编码器、掩模解码器)。

该训练过程的目标是使模型迭代地改进分割掩码,并为任何输入注释提供有效掩码输出。除了步骤4 e中的掩码采样之外,我们按照Kirilov等人的描述来实现它。在这里,原始训练方案每次都使用先前的模型预测作为下一次子迭代的掩码输入,而不是对其进行采样。我们发现,这种方法导致模型在给出多个点注释时“依赖”掩码提示的存在,如果不给出此掩码提示,则会导致性能下降。为了启用这两种设置,使用带或不带掩码提示的多点提示进行分割,我们引入了上述采样过程。使用默认SAM和通过我们的训练实现微调的模型,在有和没有掩模提示的情况下,迭代分割的定量比较见补充图1。我们还尝试了更简单的训练方案,这些方案不涉及多个子迭代,而是只从地面实况中采样框和/或固定数量的点注释。我们发现,这种方法会导致迭代分割的结果更差;该模型不能很好地进行模型预测的交互式校正。即使是更简单的训练方法,例如仅训练从MedSAM 15中完成的框提示进行分割,也会导致模型的交互分割能力进一步下降。

为了训练分割解码器(见上一节),我们将交互式分割的训练迭代和自动分割的迭代交织在一起。在这里,我们使用相同的图像和地面实况作为交互式分割的采样。我们从地面实况中推导出解码器的目标通道:中心和边界距离以及前景图(参见前面的部分和补充图2a)。然后,我们计算这些目标和解码器预测之间的损失,并根据它更新图像编码器和分割解码器的权重。我们还探索了其他两种训练策略,首先训练模型进行交互式分割,然后训练AIS,尝试更新图像编码器的权重并保持它们冻结。我们发现,训练交互式和AIS共同导致最佳结果;其他策略导致交互式分割(如果更新图像编码器权重)或AIS(如果编码器权重被冻结)的结果减少。

对于训练过程中的验证步骤,我们依赖于一个更简单的交互式分割过程,其中我们对每个对象的边界框和固定数量的点进行采样,使用地面实况和预测对象之间的平均Dice得分作为度量。对于自动分割,我们使用与训练中相同的损失函数作为度量,并将度量值相加以进行交互式和自动验证。本文中报告的所有实验都依赖于微调SAM出版物提供的权重;在一些实验中,我们进一步微调我们的模型。我们的训练方法也可以用于从随机初始化的权重训练模型。然而,我们预计这种方法会大幅增加培训时间,因此没有采用这种方法。

我们使用以下设置和超参数进行训练:

  • 我们使用的批量大小为2,也就是说,每批采样两个图像和相应的地面真值。在使用受限资源进行训练的情况下(图5a和扩展数据图10),我们使用批量大小1。进一步的训练超参数记录在扩展数据图10 c中。
  • 我们用512 × 512像素的补丁形状训练所有模型;一些训练数据集包含较小的图像,这些图像被零填充以匹配此形状。唯一的例外是LIVECell专家,我们已经用520 × 704的补丁形状(完整图像形状)训练了它。
  • 我们使用Dice损失来比较地面实况对象和掩码预测,以进行交互式和自动分割。
  • 我们使用L2损失来比较真实和预测的IOU分数。·我们使用ADAM优化器,初始学习率为10−5。我们还研究了学习率和优化器的影响,尝试了5 × 10−4、10−4、5 × 10−5、10−5和5 × 10−6的学习率,以及使用ADAM 53代替ADAM。我们发现,使用高于10−5的学习率会导致更差的结果,并且没有发现其他参数选择的影响。
  • 为了训练AIS的解码器输出,我们使用三个预测通道、中心和边界距离以及前景预测的平均Dice损失,掩盖了两个距离通道的背景损失。有些违反直觉的是,我们发现使用Dice作为距离预测的损失函数比使用L2损失更好。请注意,距离通道被归一化为范围[0,1],因此Dice损失被很好地定义。
  • 当验证指标达到平台时,我们降低学习率(ReduceLROnPlateau)
  • 对于大多数实验,我们训练模型进行 250,000 次迭代,并使用达到最佳验证指标的 epoch。我们发现,EM 和 LM 通才模型,尤其是 ViT-L 和 ViT-H,它们都是在大型和多样化的数据集上训练的,在训练后期不断改进,并将这种设置用于我们训练和比较这些模型的所有实验(图 3a 和 4a 以及扩展数据图 1b、2、5 和 8)。
  • 为了研究哪些模型部分需要微调(图1b),以及研究LIVECell上较低的训练数据部分(图1c和扩展数据图1c),我们训练模型,最多提前停止100,000次迭代。对于资源受限设置(图5b和扩展数据图10a)和用户研究(图6),我们还使用提前停止并训练最多100个epoch。
  • 除非另有说明,否则模型是在具有80 GB VRAM的A100 GPU上训练的,其中训练250,000次迭代的模型大约需要6天。我们在扩展数据图10c中提供了用于我们实验的所有硬件配置的概述,并在扩展数据图10d中列出了具有早期停止的代表性训练时间。

对于实现,我们尽可能地重用Kirilov等人的代码,并使用PyTorch 54和torch-em 55实现额外的训练逻辑,这是一个基于PyTorch的应用于显微镜的深度学习库,也是我们开发的。

Inference and evaluation

为了定量评估SAM的交互式分割,我们使用模型模拟基于用户的分割,遵循与前一节中描述的类似逻辑。

我们从地面实况对象中获得提示,使用图像和这些提示作为输入运行模型,并评估预测的掩码。我们不计算任何损失函数,也不累积梯度。我们实现了两种不同的评估方法–一种是始终使用来自前一次迭代的掩码,另一种是不使用掩码(步骤4 e;另见补充图1中的评估)。我们通过计算与地面真实掩模相比的平均分割精度来单独评估每个子迭代的结果(见下文)。在图2a、扩展数据图1b和补充图1中,我们报告了每个单独子迭代的平均分割精度,在七次迭代后停止。我们区分了从框提示(红色条)或从单点提示(绿色条)开始的情况。对于所有其他图,我们仅报告第零个子迭代(即,仅基于初始框或点提示的分割)和相应的最后一个子迭代的平均分割精度。请注意,点提示是随机采样的(受先前子迭代中的预测误差影响);我们在扩展数据图1b中研究了这种随机性的影响。

我们通过计算模型预测和地面真实掩模之间的平均分割精度(见下文)来评估自动分割模型。在评估AMG时,我们发现优化它的两个超参数也很重要:IOU和稳定性阈值,它们用于过滤掉低质量的预测(另请参见第一个方法部分)。虽然默认设置对于原始SAM模型工作良好,但对于微调模型,它们必须降低。据推测,这是因为这些模型被更好地校准到显微镜中物体的实际预测质量,与自然图像相比,这是较低的。为了有效地执行网格搜索,我们预先计算预测的对象掩码,然后评估要测试的超参数范围。在单独的验证集上执行参数搜索,并将找到的最佳设置应用于测试集。对于AIS,我们通过网格搜索确定了适用于中心和边界距离的阈值的最佳参数。在注释工具(下一节)中,AIS和AMG参数的最佳值将自动为所选模型设置。

为了与CellPose进行比较,我们使用最适合给定数据的CellPose模型(在运行实验时),对应于LIVECell和TissueNet的CellPose专家模型以及cyto 2模型。我们使用这些模型的默认设置,并使用CellPose Python库运行预测。对于MitoNet,我们使用napari插件进行2D分割,MitoNet_v1模型具有默认参数。请注意,在运行实验时,MitoNet Python库不是开源的,因此我们使用napari插件(https://github.com/volume-em/empanada-napari/)。

我们用平均分割精度来评估分割结果。分割准确度SA(t)在Everingham等人27中引入,并根据IOU阈值t下的真阳性TP(t)、假阳性FP(t)和假阴性FN(t)定义为:SA(t)= TP(t)/(TP(t)+ FP(t)+ FN(t))。TP(t)、FP(t)和FN(t)是通过在每个对象级别上匹配分割和地面实况并将具有比t更高的IOU值的匹配计数为TP(t)、将预测中的不匹配对象计数为FP(t)以及将地面实况中的不匹配对象计数为FN(t)来计算的。然后通过在0.5至0.95的范围内以0.05为增量对阈值t上的SA(t)求平均值来计算平均分割精度。我们计算每个图像的分数,然后将其平均到给定评估数据集的所有图像上。DSB细胞核分割工具26将该指标推广到显微镜检查中,Hirling等人最近在显微镜分割的背景下对该指标进行了深入研究。平均分割精度是一个严格的评估标准,因为它包括在高IOU阈值下的评估,即使是与地面实况对象的微小偏差也会受到惩罚。出于这个原因,我们在一些实验中使用了不太严格的SA(0.5)度量,我们发现它对于有意义的评估来说太严格了。

为了评估用户研究中的跟踪结果的质量,我们使用了由Cell Tracking Base 31引入的跟踪度量。该度量将由地面实况跟踪注释定义的图和由预测的跟踪结果定义的图彼此匹配,然后对该匹配中的误差进行计数。我们使用’traccuracy’存储库(https://github.com/Janelia-Trackwords-2023/traccuracy/)。

Interactive annotation tools and Python library

我们扩展了SAM的核心功能,以支持预计算的图像嵌入,平铺计算的图像嵌入和多维分割的基础上投影提示相邻的切片/时间帧的缓存。我们在μSAM Python库中实现了此功能,使用scipy 57和scikit-image 49来实现额外的图像处理逻辑。我们还使用科学Python库numpy 58,pandas 59和matplotlib 60来实现我们的库,并为本文执行额外的数据分析和绘图。我们的Python库还实现了前面几节中描述的训练和评估功能。它还包含我们的napari插件的实现,该插件实现了五个不同的小部件:用于2D注释,用于3D注释,用于高吞吐量图像注释,用于跟踪和模型训练。有关此功能实现的更详细描述,请参阅补充信息。

User study

我们进行了三个不同的用户研究,以证明我们的napari工具的2D分割,3D分割和跟踪的有用性。第一个用户研究由五个不同的注释器执行,我们比较了μSAM和CellPose在明场图像中注释类器官的情况。我们使用多个注释器进行用户研究,以研究用户之间的注释性能差异,并比较两种工具的不同注释模式。在第二个用户研究中,单个注释器在3D EM中分割细胞核,将μSAM与ilastik雕刻进行比较。在最后一项用户研究中,通过单个注释器在荧光显微镜中跟踪细胞核,比较了μSAM的交互式跟踪和TrackMate的跟踪和校正。用户研究的详细说明见补充信息。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值