CellVIT: Transformer-Based Panoptic Segmentation

• The paper introduces a Transformer-based architecture integrating a ViT encoder with a U-Net–style decoder to deliver precise cellular panoptic segmentation for digital pathology.
• It employs innovative loss functions and optimized training paradigms, including focal Tversky, Dice, and gradient-based losses, to enhance instance separation accuracy and generalization.
• The approach demonstrates robust performance on datasets like PanNuke and PUMA, achieving improved PQ, SQ, and F1 scores compared to previous methods while ensuring reproducibility and practical deployment.

CellVIT Panoptic Instance Segmentation refers to a class of Transformer-based deep learning frameworks for precise cellular panoptic segmentation in digital pathology, combining instance-level nuclei identification and type assignment with semantic tissue classification. These systems replace convolutional backbones with vision transformers (ViT), processing hematoxylin and eosin-stained (H&E) images from challenging datasets such as PanNuke and PUMA, and employ panoptic instance segmentation metrics including Panoptic Quality (PQ), Segmentation Quality (SQ), and Detection Quality (DQ) to benchmark performance. Here, "CellVIT" encompasses both the original architecture (Hörst et al., 2023) for nuclei segmentation and its extension CellViT++ (Shahamiri et al., 15 Mar 2025) as deployed in generalizable panoptic pipelines.

1. Architectural Overview

CellVIT architectures fundamentally comprise a ViT encoder and a U-Net–style decoder with skip connections, facilitating dense pixelwise and per-instance predictions for nuclei in histopathological imagery (Hörst et al., 2023). The ViT encoder embeds the input image $x \in \mathbb{R}^{H \times W \times 3}$ by partitioning it into $P \times P$ patches (commonly $P=16$) and projecting each to $D$-dimensional tokens via a linear embedding matrix. A learned $[\mathrm{CLS}]$ token is prepended, and spatial ordering is maintained through learned positional embeddings.

The encoder is scalable:

• ViT₍₂₅₆₎ (“ViT-Small”, $D=384$, $L=12$, $h=6$)
• SAM-B (ViT-Base, $D=768$, $L=12$, $P \times P$0)
• SAM-L (ViT-Large, $P \times P$1, $P \times P$2, $P \times P$3)
• SAM-H (ViT-Huge, $P \times P$4, $P \times P$5, $P \times P$6)

Encoder tokens at selected depths are routed via skip connections to a multi-stage decoder. Five skip connections (at depths $P \times P$7) are concatenated or added after upsampling and convolutional refinement at each decoder stage. Multi-task output heads generate:

• A binary nuclei mask (NP-branch)
• Horizontal/vertical (HV) distance maps for instance separation
• Per-pixel softmax cell type classification (NT-branch) into $P \times P$8 classes (PanNuke)
• Whole-image tissue type classification over 19 types, using the terminal $P \times P$9 token (TC-branch)

In CellViT++ (Shahamiri et al., 15 Mar 2025), the backbone and segmentation decoder are left unchanged, with only the final per-cell classification head fine-tuned for match to downstream data distributions (e.g., PUMA, with 3 or 10 cell taxonomy classes depending on the track). Panoptic segmentation is achieved by fusing CellViT-derived instance maps and nnU-Net–generated semantic tissue masks.

2. Loss Function Design and Training Paradigms

The total loss in CellViT is a weighted aggregate of losses from each prediction head:

$P=16$0

• NP-branch employs both Focal Tversky ($P=16$1) and Dice losses.
• HV-branch regresses distance maps and gradients using MSE and a "mean squared gradient error" term:

$P=16$2

• NT-branch (cell classification) combines Focal Tversky, Dice, and BCE losses.
• TC-branch minimizes cross-entropy on tissue classes.

Losses are hyperparameter-weighted; e.g., for optimal PanNuke results, $P=16$3. During training, encoder weights are frozen for 25 epochs before fine-tuning end-to-end, with batch size 16, an AdamW optimizer (learning rate $P=16$4, decayed by $P=16$5 per epoch), for 130 epochs (Hörst et al., 2023). For CellViT++ in PUMA, only the classifier is fine-tuned (100 hyperparameter trials), reducing compute demands (Shahamiri et al., 15 Mar 2025).

3. From Pixel Embeddings to Panoptic Instances

Nuclei separation leverages NP and HV outputs. The process is:

1. Compute Sobel gradients ($P=16$6, $P=16$7) from predicted maps.
2. Seed watershed markers at local minima of $P=16$8.
3. Execute marker-controlled watershed to segment each nucleus instance.

Post-processing assigns the majority cell type (from NT) per instance. For CellViT++ in PUMA, mask thresholding (0.5), morphological opening, and removal of objects $P=16$920 px suffice; no watershed is necessary due to transformer decoder separation capabilities. The instance label overlay produces the canonical panoptic map, integrating cell instances and semantic tissue class predictions (Shahamiri et al., 15 Mar 2025).

4. Pretraining, Augmentation, and Generalization

CellViT variants utilize expansive pretraining:

• In-domain: ViT₍₂₅₆₎ self-supervised on 104M TCGA histology patches (DINO)
• Out-of-domain: SAM pretraining (1.1B masks/11M "natural" images)

Augmentation strategies (Albumentations library) include rotational, flipping, scaling, noise, blur, elastic, and color jitter transformations. Oversampling ($D$0) remedies rare class imbalance. PUMA's CellViT++ pipeline similarly applies flips, 90° rotations, and ±10% color augmentation (Hörst et al., 2023, Shahamiri et al., 15 Mar 2025).

This regimen is shown to:

• Boost F₁ performance from 0.78 (none) to 0.82 (aug-only); "dead" class F₁ +0.13.
• Generate robust generalization to MoNuSeg (bPQ=0.672), with large-patch inference (1024×1024 px, 64 px overlap), confirming instance segmentation viability at gigapixel scale (Hörst et al., 2023).

5. Panoptic Metrics and Benchmark Results

Evaluation leverages panoptic quality (PQ), segmentation quality (SQ), detection quality (DQ), and traditional F₁/Dice:

$D$1

• TP: predicted-ground-truth mask pairs with $D$2
• F₁ (detection): center-of-mass distance $D$36 px

In PanNuke cross-validation (Hörst et al., 2023):

• CellViT-SAM-H: mPQ = 0.4980
• HoVer-Net: mPQ = 0.4629 ($D$47.7% rel.)
• STARDIST/RN50: mPQ = 0.4796 ($D$53.9% rel.)
• Detection F₁: 0.83 (CellViT-SAM-H) vs 0.80 (HoVer-Net)
• Enhanced instance separation across all cell types; "dead" class is most challenging (PQ = 0.149, F₁ = 0.43)

For PUMA (Shahamiri et al., 15 Mar 2025), panoptic segmentation combining CellViT++ and nnU-Net yields:

• Tissue Dice: 0.750 (vs baseline 0.629)
• Nuclei F₁ (Track 1): 0.611 (vs baseline 0.638)
• $D$6 ($D$7 above baseline)
• Improvements arise primarily from better semantic tissue segmentation, while detection is maintained.

6. Variant Analysis, Limitations, and Interpretability

Ablation studies (Hörst et al., 2023):

• Encoder pretraining: F₁ improves from 0.80 (random) to 0.83 (SAM-H pretr.)
• Oversampling (γ_s) and Focal Tversky loss foster rare-class and minority tissue performance.
• Decoder choice: HoVer-Net style outperforms Stardist/CPP-Net.
• Performance is highly sensitive to input resolution: F₁ drops from 0.83 (0.25 μm/px) to 0.71 (0.50 μm/px), impacting tiny and "dead" nuclei.
• Challenges remain in segmenting highly clustered nuclei and extremely small objects.
• In PUMA, CellViT++ required no region-growing due to the decoder's inherent ability to separate touching cells; post-processing only involved simple morphological filtering.

7. Reproducibility, Code, and Application Integration

CellViT is open-sourced at [https://github.com/TIO-IKIM/CellViT] under a permissive MIT-style license for non-commercial use (Hörst et al., 2023), while the PUMA pipeline code (CellViT++ + nnU-Net) is at https://github.com/TIO-IKIM/PUMA. Key dependencies include PyTorch ≥1.13.1, Albumentations, and CUDA toolkit, and the approach is designed for a single GPU (A100 80 GB or RTX A6000 48 GB). Large-patch inference strategies, efficient memory fusion of segmentations, and direct QuPath-compatible JSON export provide downstream usability. Inference can be readily adapted via: $D$8 This demonstrates practical deployment, enabling rapid, panoptic-grade segmentation and classification workflows for digital pathology.