Source-linked AI summary

CellViT: Vision Transformers for Precise Cell Segmentation and Classification

Fabian Hörst, Moritz Rempe, Lukas Heine, Constantin Seibold, Julius Keyl, Giulia Baldini, Selma Ugurel, Jens Siveke, Barbara Grünwald, Jan Egger, Jens Kleesiek

arXiv:2306.15350v2eess.IVcs.CVcs.LG

TL;DR

Nuclei detection and segmentation in H&E tissue images are clinically important but difficult because nuclei vary in staining and size, overlap, and cluster. CellViT addresses this with a Vision Transformer-based segmentation architecture using pretrained encoders, achieving state-of-the-art PanNuke performance while also supporting faster inference and cross-dataset evaluation.

  • Problem

    Automated nuclei detection and segmentation remain challenging in H&E tissue images because nuclei vary in staining and size, overlap, and cluster, while reliable cell-level analysis is needed for clinical and downstream applications.

  • Method

    CellViT uses a Vision Transformer encoder within a U-Net-shaped encoder-decoder architecture, combining pretrained ViT or SAM encoders with skip-connected upsampling for nuclei segmentation and feature extraction.

  • Results

    CellViT achieves state-of-the-art nuclei detection and instance-segmentation performance on PanNuke, including a mean panoptic quality of 0.50 and an F1-detection score of 0.83.

  • Takeaways & Limitations

    CellViT provides a single architecture for nuclei detection, segmentation, and simultaneous nuclei feature extraction, with application to MoNuSeg without finetuning and faster inference than HoVer-Net.

  • Takeaways & Limitations

    Dead-cell instance segmentation remains difficult because dead nuclei are small, and downscaling images from 0.25 µm/px to 0.50 µm/px makes detection and segmentation harder.

Abstract

from arXiv · show

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

1 Introduction

CellViT introduces a Vision Transformer-based architecture for automated nuclei instance segmentation and feature extraction in digitized tissue images. It combines pretrained Transformer encoders with decoder skip connections and reports strong PanNuke performance and faster inference than HoVer-Net.

  • CellViT addresses the need for automated, reliable cell detection, segmentation, and feature extraction in large-scale whole-slide image analysis.The motivation includes the time cost and observer variability of manual cell-level analysis.
  • The model combines a Transformer encoder with multiple skip connections and a dedicated upsampling decoder to derive precise instance segmentations.Nuclei embeddings are extracted from the Transformer encoder.
  • CellViT is a Vision Transformer-based architecture for automated instance segmentation of cell nuclei in digitized tissue samples.
  • Vision Transformers capture local and global context through attention and provide localizable cell feature vectors during the same forward pass as detection.At common image resolutions, 16-pixel tokens are approximately cell-sized, enabling direct association with cell detections.
  • CellViT surpasses existing nuclei detection methods by a substantial margin, achieves competitive PanNuke segmentation results, and generalizes to MoNuSeg without finetuning.
  • 1.85 times faster inference than HoVer-Net is achieved using 1024×1024-pixel patches instead of conventional 256-pixel patches.

2 Related Work

Prior nuclei instance-segmentation methods largely rely on handcrafted features or CNN-based architectures, while Transformer-based encoder-decoder designs provide an alternative for modeling broader image context. Large-scale pretraining is especially relevant because Vision Transformers require substantial data to learn meaningful representations.

  • Traditional nuclei-segmentation methods use handcrafted intensity, texture, shape, and morphological features, whereas deep learning automatically extracts task-relevant features.
  • Two-stage deep-learning methods detect nuclei before refining segmentations, but overlapping neighbors can require additional cleanup and some architectures are computationally complex.
  • Boundary-based methods directly predict nuclear contours, with architectures such as DCAN using U-Net and TSFD-Net using a multi-scale Feature Pyramid Network.
  • Existing promising nuclei instance-segmentation models are predominantly CNN-based, although CNNs may struggle to capture spatial long-range relationships because of local receptive fields.
  • Transformer instance-segmentation approaches commonly integrate Transformer encoders into U-Net-like or other encoder-decoder architectures, with variations in tokenization, decoders, and skip connections.
  • Vision Transformers generally need more data than CNNs to learn meaningful representations, making large-scale pretraining important for their use on smaller medical datasets.
  • The ViT256 encoder was pretrained with DINO on 104 million histological image patches and is particularly relevant to semantic segmentation.

3 Methods

CellViT adapts a ViT encoder-decoder architecture for 2D nuclei instance segmentation, classification, and feature extraction. Multiple decoder branches and postprocessing combine complementary predictions to separate overlapping nuclei and assign cell types.

  • Multi-task Decoder: Three output branches predict binary nuclei maps, horizontal-vertical distance maps, and nuclei-type instance maps.The branches correspond to nuclei prediction, horizontal-vertical prediction, and nuclei-type prediction.
  • Postprocessing: Postprocessing merges branch outputs, separates overlapping nuclei, and determines each nucleus class from the nuclei-type map.The separation pipeline uses gradients, Sobel edge detection, and marker-controlled watershed.
  • Network Structure: CellViT replaces a CNN encoder with a Vision Transformer connected to an upsampling decoder through skip connections.The design preserves access to fine-grained information while using the original ViT structure with pre-trained ViT256 and SAM encoders.
  • Transformer Encoder: The Transformer processes flattened image patches as token embeddings, with token count N = HW/P^2 and projection into a D-dimensional latent space.Alternating multiheaded self-attention and multilayer perceptrons form the Transformer blocks, while decoder stages fuse multi-scale features through skip connections.
  • Target and Losses: The total training loss combines branch-specific losses for nuclei prediction, horizontal-vertical prediction, nuclei typing, and tissue classification.Branch contributions are controlled by hyperparameters, with segmentation, distance-map, gradient, and cross-entropy terms used across tasks.
  • Cell Embeddings: Cell embeddings are extracted during the forward pass and spatially linked to detected nuclei, avoiding an additional cropped-cell inference pass.These embeddings are intended as cell features for downstream disease prediction, treatment response, and survival prediction.

4 Experimental Setup

The experiments train and evaluate CellViT primarily on PanNuke, with additional external datasets and analyses covering segmentation, detection, embeddings, and inference. Evaluation uses instance-aware metrics and experiments address architecture, imbalance, augmentation, and input handling.

  • Datasets: PanNuke contains 189,744 annotated nuclei in 7,904 256×256 px images spanning 19 tissue types and 5 cell categories.Images were acquired at ×40 magnification with 0.25 µm/px resolution, and the dataset is highly imbalanced across tissues and nuclei classes.
  • Datasets: MoNuSeg provides an additional segmentation evaluation set of 14 1000 × 1000 px images containing more than 7000 nuclei across seven organ types.Unlike PanNuke, MoNuSeg does not divide nuclei into different classes, and only its test set is used.
  • Datasets: CoNSeP is used as an external validation dataset for analyzing embeddings from detected cells in colorectal adenocarcinoma whole-slide images.The dataset comprises 41 H&E-stained whole-slide images, rescaled to 1024 × 1024 px for analysis.
  • Experimental Design: Experiments compare pretrained and randomly initialized models, decoder designs, regularization choices, and customized oversampling before evaluating selected models for segmentation quality.The study includes two PanNuke experiments, one MoNuSeg experiment, and an internal inference-speed comparison.
  • Evaluation Metrics: Panoptic quality is used because Dice and Jaccard do not evaluate nuclei detection, separation of overlapping instances, and instance segmentation together.Binary PQ treats all nuclei as one class, while multi-class PQ accounts for nuclei categories and decomposes into detection and segmentation quality.
  • Evaluation Metrics: Detection evaluation reports precision, recall, and F1 using true positives, false positives, and false negatives across nuclei classes.For detection matching, predicted and ground-truth nuclei are paired when their centers of mass lie within a 6 px radius.
  • Training Strategy: Data augmentation combines geometric, noise, intensity, superpixel, blur, cropping, and elastic-transformation operations to increase variety and discourage overfitting.The augmentation policy includes random rotations and flips, downscaling, Gaussian noise, color jittering, and random cropping with resizing.

5 Results

Experiments evaluate CellViT for nuclei detection, segmentation, cross-dataset generalization, token-based analysis, and inference speed. Pre-trained encoders, large input tiles, and integrated embeddings provide strong performance, while dead-cell segmentation and resolution shifts remain challenging.

  • Detection Quality on PanNuke: Pre-trained CellViT encoders substantially improve detection, with CellViT-SAM-H performing best and epithelial-nuclei F1 increasing by up to 26%.ViT256 and SAM encoders perform at a similar level and outperform baseline and detection-only methods.
  • Segmentation Quality on PanNuke: CellViT256 and CellViT-SAM-H achieve strong panoptic quality across neoplastic, connective, and epithelial nuclei, while dead cells remain the weakest category.Oversampling improves most classes but can reduce neoplastic-nuclei quality; removing Focal Tversky loss decreases PQ across classes.
  • MoNuSeg Test Performance: Dead-cell instance segmentation becomes more difficult after downscaling images from 0.25 µm/px to 0.50 µm/px.The difficulty is attributed to dead cells’ small size.
  • MoNuSeg Test Performance: 0.672 bPQ is achieved by CellViT-SAM-H on 1024 px MoNuSeg tiles without patching, comparable to 0.671 with overlapping 256 px patches.Non-overlapping 256 px patches reduce bPQ to 0.631, while larger tiles do not degrade performance.
  • Token Analysis: CellViT-SAM-H embeddings classify 8,773 CoNSeP nuclei with AUROC 0.963, while CellViT256 embeddings reach AUROC 0.960.The embeddings are extracted during the same inference pass as cell detection and support distinctions among nuclei classes and tissue environments.
  • Inference Runtime: 2.49× and 2.25× speedups are obtained for CellViT256 and CellViT-SAM-H, respectively, with 1024 px rather than 256 px input patches.Using 1024 px patches, CellViT256 and CellViT-SAM-H are 1.85× and 1.39× faster than HoVer-Net.

6 Discussion and Conclusion

CellViT achieves state-of-the-art nuclei instance segmentation and detection on PanNuke, while results on MoNuSeg support generalizability to unseen cohorts.

  • CellViT demonstrates state-of-the-art performance for nuclei instance segmentation and detection on PanNuke.
  • CellViT generalizes to the MoNuSeg dataset without finetuning, validating performance on previously unseen cohorts.
  • Nuclei detection and classification surpass existing methods by a significant margin, while accurate detection is prioritized for large-scale clinical analysis.

CRediT authorship contribution statement

The authors’ contributions span methodology, software, validation, writing, supervision, resources, and project administration.

  • Fabian Hörst contributed across conceptualization, methodology, software, analysis, investigation, data curation, writing, and visualization.
  • Moritz Rempe and Lukas Heine contributed methodology and writing, while Constantin Seibold contributed conceptualization and editing.
  • Julius Keyl and Giulia Baldini contributed validation and writing review, with additional supervision and project roles covered by other authors.
  • Jens Kleesiek contributed resources, supervision, project administration, funding acquisition, and writing review.

Declaration of competing interest

The declaration reports external research support, honoraria, consultancy, educational presentations, and meeting or travel support for named authors.

  • S. U. reports research support from Bristol Myers Squibb and Merck Serono outside the submitted work.
  • S. U. reports speakers’ and advisory-board honoraria from multiple pharmaceutical companies.
  • J.T.S. receives consultancy and continuing-medical-education honoraria from listed pharmaceutical companies.

Data availability

The datasets are publicly available, and the CellViT models and source code are available online.

  • The datasets used in the study are publicly available.
  • All models and source code are available through the CellViT GitHub repository.

A.1 STARDIST and CPP-Net

This section describes modified loss functions for STARDIST and CPP-Net, including their component branches and weighting factors.

  • STARDIST: STARDIST’s loss function changes because of the probability branch PD and radial distances branch RD.The section introduces these branches as the reason for the revised objective.
  • Loss branches: The individual loss branches include BCE, MSE, and Dice-based terms for probability, spatial-distance, and nuclei-type objectives.The listed terms are expressed using the branch-specific weighting factors and loss components.
  • Weighting factors: The reported weighting factors are λPDBCE = λSDMSE = λNTDICE = λNTBCE = 1.These coefficients assign unit weight to the stated loss components.
  • CPP-Net: CPP-Net uses the same loss structure as equation (12) but replaces the nuclei-type focal term with λNTFT = 0.5, λNTDICE = 0.2, and λNTBCE = 0.5.The passage states that this setting achieved superior results for CPP-Net.

A.2 Supplementary Tables

The supplementary tables document evaluation settings, augmentation choices, and training hyperparameters for the compared nuclei-segmentation networks.

  • Evaluation: Table A.1 reports average mPQ and bPQ across 19 PanNuke tissue types using three-fold cross-validation.It also provides split standard deviations and compares downscaled 0.50 µm/px with original 0.25 µm/px images.
  • Evaluation: The CellViT results in Table A.1 use the HoVer-Net decoder, while the first two columns contain models trained and evaluated on original-resolution images.Models trained on downscaled images are marked with an asterisk.
  • Training setup: Table A.2 lists selected Albumentations data-augmentation techniques, their probabilities, and additional hyperparameters.STARDIST uses only spatial transformations, while CPP-Net adopts the HoVer-Net decoder augmentations.
  • Training setup: Table A.3 records CellViT hyperparameters for PanNuke training runs, including loss weights and focal-term parameters.The listed settings include λHVMSE = 2.5, λHVMSGE = 8, and αFT = 0.7, βFT = 0.3, γFT = 4/3, εFT = 1 · 10^-6.
  • Training setup: Table A.4 provides the CPP-Net hyperparameters used for all PanNuke training runs.The passage identifies the table’s scope but does not enumerate its parameter values.

A.3 Supplementary Figures

The supplementary figures illustrate CellViT-SAM-H predictions across imaging conditions and demonstrate inference-pipeline interoperability with QuPath on whole-slide images.

  • Prediction examples: Figure A.1 shows a MoNuSeg tissue sample with ground-truth binary masks and CellViT-SAM-H predictions.The predictions are displayed for different input sizes and magnifications.
  • Tissue examples: The supplementary material includes an esophageal adenocarcinoma tissue slide acquired at a stated magnification.The supplied passage does not provide the magnification value.
  • Pipeline interoperability: Figure A.2 presents exemplary whole-slide images with corresponding cell polygons imported into QuPath.The visualization focuses on interoperability between the inference pipeline and QuPath.
  • Pipeline interoperability: Approximately 150,000 nuclei were detected for each Figure A.2 whole-slide image and imported without fast-loading or zooming problems on a standard laptop.The first slide was acquired at ×40 and 0.25 µm/px; the second at ×20 and 0.50 µm/px.
Loading 2306.15350v2…