DIR CRC Morphology-Based Cell-Type Model (Mask2Former)

Instance segmentation + typing of individual cells in H&E-stained colorectal cancer (CRC) whole-slide images, from H&E morphology alone (no IHC/multiplex staining required). Part of the Digital Immune Reporter (DIR) pipeline.

  • Architecture: Mask2Former (Swin-S backbone), via MMDetection
  • Task: 6-class cell instance segmentation on H&E patches
  • Classes: Neutrophil, Tumor, Lymphocyte, Eosinophil, Plasmacell, Other (plus an internal Background class at index 0, not emitted as a predicted instance)
  • Input: H&E image tiles

Files

  • epoch_36.pth β€” model checkpoint (MMEngine format; includes meta.dataset_meta with the class-name list used at training time)
  • mask2former_swin-s-3x_dataset_tcga_lizard_class_weight_log_count.py β€” full MMDetection training config

Usage

from mmdet.apis import init_detector, inference_detector

model = init_detector(
    "mask2former_swin-s-3x_dataset_tcga_lizard_class_weight_log_count.py",
    "epoch_36.pth",
    device="cuda:0",
)
result = inference_detector(model, "path/to/tile.png")

model.dataset_meta["classes"] gives the ordered class-name tuple corresponding to each predicted instance's label index.

Training data

Trained on a combined dataset from two sources:

  • TCGA-CRC: tumor-rich ROIs selected by a pathologist, with cell types classified using QuPath cell classifiers. Split at the slide level into 469 training (TCGA-train) and 109 held-out test (TCGA-test) ROIs β€” the same train/test split used for the tissue-compartment segmentation model.
  • Lizard (public): 190 training and 48 held-out test images with per-cell annotations across five tissue sources. Lizard's connective-tissue class was relabeled Other to align with the TCGA annotation schema.

All images were processed at 0.5 mpp using 256x256 px patches (128x128 um).

An architecture-selection phase (Mask2Former and Mask R-CNN, across ResNet-50/101, Swin-T, Swin-S, and a UNI-adapter backbone, 12 epochs each) was run on TCGA data alone, evaluated on TCGA-test with COCO instance-segmentation metrics (bbox/segm mAP@0.50 and mAP@0.75). Mask2Former with a Swin-Small backbone achieved the best overall performance and was selected as the final architecture, then retrained from scratch on the combined TCGA-train + Lizard-train sets for 36 epochs (MMDetection v3.3.0) β€” this checkpoint.

Training used AdamW (lr 1e-4, weight decay 0.05, beta=(0.9, 0.999)) with a multi-step schedule (decay x0.1 at epochs 8 and 11), gradient clipping (max L2 norm 0.01), and a reduced backbone learning-rate multiplier (0.1). Data augmentation included random channel shuffling/shifting, random tone curve (scale 0.4) and gamma (40-60) adjustment, color jitter (hue +/-0.1, saturation +/-0.5), random cropping (50% of image dimensions), and random horizontal flip (p=0.5).

To address class imbalance, the classification loss was weighted per class using log-transformed, inverted, and normalized TCGA-train cell frequencies: 1.29 (neutrophil), 0.90 (tumor), 1.08 (lymphocyte), 1.36 (plasma cell), 1.44 (eosinophil), 0.93 (other); the background/no-object class was weighted 0.1.

Evaluated on held-out TCGA-test (N=109 ROIs) and Lizard-test (N=48 images), matching predicted and reference cell centroids via one-to-one Hungarian assignment (max radius 15 px, after class-agnostic NMS at IoU 0.5 and iterative polygon-overlap cleaning), with detection reported as macro-averaged precision/recall/F1 and classification accuracy computed only on correctly localized cells.

License

CC-BY-NC-4.0 β€” research/non-commercial use only.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support