[论文解读] CellViT: Vision Transformers for Precise Cell Segmentation and Classification

Fabian Hörst, Moritz Rempe+9 more|arXiv (Cornell University)|Jun 27, 2023

AI in cancer detection被引用 19

一句话总结

CellViT 引入基于 Vision Transformer 的类似 U-Net 的架构用于 PanNuke 的细胞核实例分割和分类，利用 ViT 编码器和多任务解码实现最先进的指标。

ABSTRACT

Nuclei detection and segmentation in hematoxylin and eosin-stained (H&E) tissue images are important clinical tasks and crucial for a wide range of applications. However, it is a challenging task due to nuclei variances in staining and size, overlapping boundaries, and nuclei clustering. While convolutional neural networks have been extensively used for this task, we explore the potential of Transformer-based networks in this domain. Therefore, we introduce a new method for automated instance segmentation of cell nuclei in digitized tissue samples using a deep learning architecture based on Vision Transformer called CellViT. CellViT is trained and evaluated on the PanNuke dataset, which is one of the most challenging nuclei instance segmentation datasets, consisting of nearly 200,000 annotated Nuclei into 5 clinically important classes in 19 tissue types. We demonstrate the superiority of large-scale in-domain and out-of-domain pre-trained Vision Transformers by leveraging the recently published Segment Anything Model and a ViT-encoder pre-trained on 104 million histological image patches - achieving state-of-the-art nuclei detection and instance segmentation performance on the PanNuke dataset with a mean panoptic quality of 0.50 and an F1-detection score of 0.83. The code is publicly available at https://github.com/TIO-IKIM/CellViT

研究动机与目标

• 在数字化的 H&E 切片中激发对细胞核的精准分析，以支持诊断和下游生物标志物发现。
• 开发一个利用 Vision Transformers 的 U-Net 风格框架的细胞核实例分割模型。
• 利用大规模预训练的 Vision Transformers 和基础模型（ViT256 与 SAM）以改进分割与分类。
• 提供一个使用大尺度补丁的用于千兆像素级全切片图像（WSIs）的高效推断管线。
• 实现细胞核类型分类和组织级引导，以改进编码器学习。

提出的方法

• 将输入图像转换为 16x16 像素的令牌序列用于 ViT 编码，带可学习的类别令牌和二维位置嵌入。
• 采用一个 ViT 编码器，通过五个跳跃连接与多分支 U-Net 风格解码器相连，具有 NP（nuclei，细胞核）、HV（水平/垂直距离）和 NT（nuclei type，细胞核类型）分支。
• 引入组织分类分支（TC），使用 PanNuke 的组织标签来引导编码器学习。
• 使用组合损失对 NP、HV、NT 和 TC 分支进行训练，采用 BCE、DICE、MSE、MSGE 和 Focal Tversky 损失，如公式所述。
• 应用后处理（受 HoVer-Net 启发）以分离重叠的细胞核，并通过对 NT 映射的多数投票来分配细胞核类型。
• 评估在 104M 组织学补丁上预训练的 ViT256（ViT256）以及 Segment Anything Model（SAM）作为骨干编码器，进行迁移学习以及域内/域外预训练。

实验结果

研究问题

• RQ1 在组织学数据上预训练的 Vision Transformers 是否能比基于 CNN 的方法在 PanNuke 上改进细胞核实例分割？
• RQ2 将 ViT 编码器与多分支解码器集成如何影响跨组织类型的分割精度与细胞核分类？
• RQ3 预训练选择（ViT256、SAM）对 PanNuke 的性能以及对 MoNuSeg 等其他数据集的泛化有何影响？
• RQ4 快速的大补丁推断策略（1024x1024 px）在保持精度的同时，是否支持 WSI 级分析？
• RQ5 组织级监督是否能引导编码器表征以改进下游的细胞核分类？

主要发现

• 在 PanNuke 上实现了最先进的细胞核检测与实例分割，平均全景质量 PQ 0.50，F1-detection 0.83。
• 展示了在域内和域外的大规模预训练 Vision Transformers（ViT256 和 SAM）在细胞核分析中的有效性。
• 无需微调即可泛化到 MoNuSeg，表明跨数据集的鲁棒性。
• 通过 1024x1024 px 补丁实现更快的 WSI 推断，速度比 HoVer-Net 基线快 1.85 倍，同时保持分割质量。
• 提供一个统一的框架，在单次前向传播中获得细胞核检测、定位和类型分类，相比仅 CNN 的流程减少运行时。