mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-012345

SegTME-UNI2-UperHoVer — Stage 3 (TCGA-UT Scales 0–5, Final Model)

Stage 3 of 3 in the SegTME-UNI2 progressive pseudo-label curriculum — the final deployed model. Trained from scratch on entropy-filtered pseudo-labels generated by the Stage 2 model (M2, TCGA-UT Scale 0) applied to all 1,608,061 TCGA-UT patches across six resolution scales (0.5–1.0 µm/px).

Part of the SegTME-UNI2 framework — Segmentation of Tumour Microenvironment with UNI2, targeting end-to-end cell segmentation, TME feature extraction, and clinical narrative generation from routine H&E histology. Submitted to Medical Image Analysis (Elsevier).

---

Framework Overview

SegTME-UNI2 addresses the core bottleneck in computational TME analysis: the gap between the scale of available TCGA H&E image data (1,608,060 patches) and the scale of human-annotated pixel-level labels (7,901 PanNuke images, 189,744 annotated nuclei). The solution is a three-stage progressive pseudo-label curriculum that closes this gap without additional manual annotation, staging domain expansion across resolution scales.

Stage 1: PanNuke (7,901 patches, human GT) → M1
     ↓ M1 infers on TCGA-UT Scale 0 → entropy-filtered pseudo-labels
Stage 2: TCGA-UT Scale 0 (271,711 patches, pseudo-labels) → M2
     ↓ M2 infers on TCGA-UT Scales 0–5 → entropy-filtered pseudo-labels
Stage 3: TCGA-UT Scales 0–5 (1,608,061 patches, pseudo-labels) → M3  ← this model

Critical design principle: Each stage trains a completely independent model from scratch (UNI2-h pretrained backbone + randomly initialised decoders). No weights are transferred between stages — improvement is driven entirely by increasing pseudo-label quality.

Why M2 (not M1) generates Stage 3 pseudo-labels? M2 was trained on scale-0 TCGA-UT data, so it produces substantially higher-fidelity predictions for scales 1–5 (coarser resolutions not seen in PanNuke) than M1, which only saw PanNuke at 0.25 µm/px.

Pseudo-label generation: M2 runs inference on all 1,608,061 patches across six scales. Per-pixel entropy H(p) = −Σ p̂_c·log p̂_c is computed; patches are ranked by mean foreground entropy H̄. Only the top-70% most confident patches (H̄ < τ, τ = 70th percentile, recomputed independently for Stage 3) are retained as training data.

---

Training Data — TCGA-UT Six-Scale Corpus

Scale	Patches	Resolution (µm/px)	Physical area	Curriculum role
0	271,711	0.5	128 × 128 µm	Stages 2 and 3
1	271,090	0.6	153.6 × 153.6 µm	Stage 3 only
2	269,880	0.7	179.2 × 179.2 µm	Stage 3 only
3	268,120	0.8	204.8 × 204.8 µm	Stage 3 only
4	265,460	0.9	230.4 × 230.4 µm	Stage 3 only
5	261,800	1.0	256 × 256 µm	Stage 3 only
Total	1,608,061	0.5–1.0

Source: TCGA-UT Restructured (Komura et al. 2022) — 32 solid cancer types, 7,175 patients, 8,736 slides.

---

Architecture — UNI2-UperHoVer

A novel dual-head segmentation model:

Input (3 × 224 × 224)
     ↓
UNI2-h ViT-Giant backbone (pretrained on 100M+ histopathology tiles from 100,000 slides)
     ↓ multi-scale FPN taps at transformer blocks 5, 11, 17, 23
1×1 projection layers → feature pyramid at 4 scales
     ↓
UperNet Decoder (PPM + FPN fusion)
     ├─→ Semantic Head (6-class)        → semantic segmentation map
     └─→ HV Regression Head (2-channel) → horizontal/vertical gradient maps
                                            → watershed-based nuclear instance separation

UNI2-h backbone specs:

• Architecture: ViT-Giant (patch size 14 px, embedding dim d = 1536, 24 transformer blocks, 24 attention heads, SwiGLU-packed MLP, 8 register tokens)
• Pretrained on: >100 million tiles from 100,000 whole-slide images, 40+ cancer types
• Feature pyramid: blocks {5, 11, 17, 23} → channels {256, 512, 1024, 2048} at strides {s/4, s/8, s/16, s/32}

Dual decoder heads (parameters not shared between heads):

• Semantic head: num_labels=6, hidden_size=768 → 6-class per-pixel output
• HV regression head: num_labels=2, hidden_size=768 → horizontal (ch 0) and vertical (ch 1) gradient maps

Loss function:

L_total = L_sem + λ · L_hv

L_sem  = cross-entropy over valid pixels (ignore_index=255)
L_hv   = L_MSE + 2 · L_MSGE   (foreground pixels only)
λ      = 1.0

Dynamic HV target synthesis: HV maps are generated on-the-fly from semantic labels using connected-component labelling and per-nucleus centroid computation — no instance-level annotations are required at any stage.

---

Cell Classes (PanNuke Ontology)

Class	Label	Biological meaning
0	Background	Non-cellular tissue / void
1	Neoplastic	Tumour nuclei
2	Inflammatory	Immune cells (lymphocytes, neutrophils, etc.)
3	Connective	Stromal / connective tissue nuclei
4	Dead	Necrotic / apoptotic nuclei
5	Non-neoplastic Epithelial	Normal epithelial nuclei

---

Training Configuration — Stage 3

Hyperparameter	Value
Training dataset	TCGA-UT Scales 0–5 — 1,608,061 patches, 32 solid cancer types, 0.5–1.0 µm/px
Labels	Entropy-filtered pseudo-labels from Stage 2 (M2, TCGA-UT Scale 0 model)
Training split	80% train, 20% pseudo-label validation (PL-val)
Input resolution	256×256 px → resized to 224×224 for backbone
Backbone init	UNI2-h pretrained checkpoint — fresh init, no weights from Stage 2
Decoder init	Randomly initialised (both heads)
Optimiser	AdamW (β₁=0.9, β₂=0.999, weight_decay=1×10⁻²)
Learning rate	5×10⁻⁵ — linear decay: LR(t) = 5×10⁻⁵ × (T−t)/T (no warmup)
Per-device batch size	8
Gradient accumulation	8 steps (effective batch = 64 per GPU)
Number of GPUs	6 × NVIDIA A100 (DDP)
Training epochs	100
Total optimizer steps	335,100
Mixed precision	bfloat16
Compilation	torch.compile (inductor backend)
Eval frequency	every 500 steps
Checkpoint metric	Validation mean IoU (↑)
Entropy threshold τ	70th percentile of per-patch mean entropy H̄ (recomputed for Stage 3)

Note on step count: Stage 3 accumulates 335,100 total optimizer steps — substantially more than Stage 1 (24,651) or Stage 2 (212,500) despite only 100 epochs, because each epoch processes all 1,608,061 patches.

Augmentation (each with p=0.5): colour jitter (brightness/contrast/saturation ±20%, hue ±5%), HLS-space multiplicative perturbation ∈[0.9, 1.1], horizontal flip, vertical flip.

---

Results

Checkpoint	Epoch	Step	Val mean IoU	Eval loss
	~1.5	5,000	0.535	0.214
	~15	50,000	0.660	0.167
	~30	100,000	0.704	0.153
	~50	168,000	0.741	0.138
	~75	250,000	0.763	0.127
best	~100	335,000	0.7724	0.120

Evaluation protocol: mIoU values in this table are the training-time checkpoint-selection mIoU: macro-averaged Jaccard across all 6 classes (Background included) on the TCGA-UT all-scales 20% pseudo-label validation split (PL-val). This measures self-consistency on model-generated labels, not accuracy against human ground truth.

Important: The mIoU curve was still improving linearly at the 100-epoch cutoff. This model has the most potential for improvement with further training among the three stages. The higher final training loss (0.120 vs. 0.061 for M2) reflects greater label noise and domain diversity introduced by five additional resolution scales not seen in earlier stages.

The initial mIoU of 0.53 (vs. 0.47 for Stage 2) demonstrates that M2's pseudo-labels encode higher-quality TCGA-domain information than M1's, providing a better starting point for M3.

Paper Table 7 values (5-class, Background excluded):

• Part A (cross-domain, PanNuke 20% held-out GT): M3 = 0.1587 — this collapse reflects a systematic magnification mismatch in the Stage 3 pseudo-label targets (an internal training artefact unrelated to input resolution); a corrected Stage 3 retraining is planned.
• Part B (PL-val self-consistency, TCGA-UT all-scales held-out): M3 = 0.4425 (compare: M2 on TCGA-UT s0 PL-val = 0.3666, confirming M3 achieves better self-consistency across the full six-scale distribution).

---

Qualitative Validation — HNSCC Case Studies

Note: The four-case HNSCC qualitative validation in the paper was conducted using Stage 2 (M2), not Stage 3. M3 was not used for the case study because it currently exhibits lower segmentation quality due to a systematic magnification mismatch in its Stage 3 pseudo-label targets — an internal training artefact unrelated to input resolution, diagnosed and documented in the paper's Discussion section. A corrected Stage 3 retraining is planned. See Stage 2 model card for the case study details.

The qualitative HNSCC case study used the M2 checkpoint applied to four independent head and neck squamous cell carcinoma patches from the Nasopharyngeal Research Group, UM Cancer Research Institute (UMCRI), Universiti Malaya (not in TCGA-UT training corpus).

Cases represent diverse TME phenotypes and staining conditions:

Case	Staining	TME Phenotype
A	Macenko stain-normalised (lavender-violet)	Proximal · High Mixing · Moderate desmoplasia
B	Deep magenta	Immune Hot · High Purity · High Mixing
C	Pale lavender-pink	Segregated · Compartmentalised
D	Pale pink (desmoplastic)	Immune Excluded · Stroma-Dominant · Desmoplastic

M2 produces biologically coherent nucleus outlines and class assignments across all four cases without case-specific fine-tuning. Notably, even in Case D (pale pink, reduced nuclear-cytoplasm contrast), aggregate TME statistics remain clinically representative — illustrating that TME characterisation accuracy is decoupled from individual cell segmentation precision.

---

Inference

# MPP normalisation before inference
scale = mpp_input / 0.314   # target ~0.35 µm/px

# Tiling
# Tile 224×224 with 50% overlap (stride 112 px); zero-pad to multiple of 14
# Stitch semantic map (foreground-priority) and HV map (last-write)

# Watershed instance separation
# Energy = HV magnitude; EDT seeds; compactness=0.01

The model expects 224×224 RGB patches normalised with ImageNet statistics (μ=(0.485, 0.456, 0.406), σ=(0.229, 0.224, 0.225)).

---

Repository Contents

File	Description
model.safetensors	Model weights — UNI2-h backbone + both decoder heads (~3.2 GB)
trainer_state.json	Full training log: per-step mIoU, loss, learning rate
training_args.bin	HuggingFace Trainer configuration

---

Related Models in This Series

Model	Stage	Training data	Best val mIoU	Eval basis
PanNuke	1	PanNuke — 7,901 patches (GT)	0.9313	Human GT
TCGA-UT-0	2	TCGA-UT Scale 0 — 271K patches (PL)	0.8197	PL-val†
This model	3	TCGA-UT Scales 0–5 — 1.6M patches (PL)	0.7724	PL-val†

†PL-val = pseudo-label validation (self-consistency on model-generated labels, not human GT). Cross-domain evaluation on PanNuke GT for M2/M3 is pending (requires dedicated inference run).

---

Data Availability

• TCGA data: https://portal.gdc.cancer.gov/
• PanNuke dataset: https://warwick.ac.uk/fac/cross_fac/tia/data/pannuke
• UNI2-h backbone: https://huggingface.co/MahmoodLab/uni2-h
• Pseudo-labelled TCGA-UT dataset: https://huggingface.co/datasets/mizjaggy18/tcga-ut-cell-instance-semantic
• Training code: to be released upon paper acceptance

---

Citation

This model is part of the SegTME-UNI2 framework (manuscript under review). If you use this work, please cite the SegTME-UNI2 paper. Please also cite the UNI2 foundation model, PanNuke dataset, and TCGA-UT dataset:

@article{wanahmad2026segtme,
  title={SegTME-UNI2: A Foundation Model-Based Framework for Generalisable Multiclass Cell Segmentation and LLM-Driven Tumour Microenvironment Characterisation in Histopathology},
  author={Wan Ahmad, Wan Siti Halimatul Munirah and Samidi, Faris Syahmi and Ahmmed, Mohammad Badal and Thiviyanathan, Vimal Angela and Thavaraj, Selvam James and Abdul Majeed, Anwar P.P.},
  journal={arXiv preprint arXiv:2606.17702},
  year={2026},
  doi={10.48550/arXiv.2606.17702}
}
@article{chen2024uni2,
  title={UNI2: Towards a Universal Whole-Slide Foundation Model for Pathology},
  author={Chen, Richard J and others},
  journal={arXiv preprint arXiv:2406.01647},
  year={2024},
  doi={10.48550/arXiv.2406.01647}
}
@article{komura2022tcgaut,
  title={Universal encoding of pan-cancer histology by deep texture representations},
  author={Komura, Daisuke and others},
  journal={Cell Reports},
  volume={38},
  pages={110424},
  year={2022}
}
@article{gamper2020pannuke,
  title={PanNuke Dataset Extension, Insights and Baselines},
  author={Gamper, Jevgenij and others},
  year={2020}
}