Metadata-Version: 2.4
Name: segtme-uni2
Version: 0.3.0
Summary: SegTME: UNI2 + UperHoVer dual-head nuclei/tissue segmentation
Project-URL: Paper, https://arxiv.org/abs/2606.17702
Project-URL: HuggingFace, https://huggingface.co/mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-0
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0
Requires-Dist: transformers>=4.35
Requires-Dist: timm>=0.9
Requires-Dist: safetensors>=0.4
Requires-Dist: huggingface_hub>=0.20
Provides-Extra: predict
Requires-Dist: numpy>=1.21; extra == "predict"
Requires-Dist: pillow>=9.0; extra == "predict"
Requires-Dist: torchvision>=0.15; extra == "predict"
Requires-Dist: opencv-python>=4.5; extra == "predict"
Requires-Dist: scipy>=1.7; extra == "predict"
Requires-Dist: scikit-image>=0.19; extra == "predict"
Dynamic: description
Dynamic: description-content-type
Dynamic: project-url
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# segtme-uni2

**SegTME** — Tumor Microenvironment segmentation using the [UNI2](https://huggingface.co/MahmoodLab/UNI2-h) pathology foundation model with a dual-head UperNet + HoverNet-style decoder.

The model simultaneously predicts:
- **Semantic segmentation** — 6-class tissue map (background, neoplastic, inflammatory, connective, dead, epithelial)
- **HV maps** — horizontal/vertical nuclear distance gradients for instance segmentation via watershed post-processing

Architecture and model weights are separate: weights are publicly hosted on HuggingFace; the architecture is distributed as this compiled package (code not exposed).

---

## Available Models

| Model | HuggingFace Repo | Trained on | Best mIoU |
|---|---|---|---|
| M1 — PanNuke | [SegTME-UNI2-UperHoVer_PanNuke](https://huggingface.co/mizjaggy18/SegTME-UNI2-UperHoVer_PanNuke) | PanNuke (pan-cancer nuclei, 7,901 patches) | **0.9313** |
| M2 — TCGA-UT-0 | [SegTME-UNI2-UperHoVer_TCGA-UT-0](https://huggingface.co/mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-0) | TCGA-UT subset 0 (~850 steps/epoch) | **0.8197** |
| M3 — TCGA-UT-012345 | [SegTME-UNI2-UperHoVer_TCGA-UT-012345](https://huggingface.co/mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-012345) | TCGA-UT subsets 0–5 combined | **0.7724** |

All three models share the same architecture and hyperparameters; only training data and checkpoint weights differ.

---

## Architecture

```
Input (B, 3, 224, 224) — ImageNet-normalised
       │
       ▼
UNI2 ViT-Giant backbone (depth=24, heads=24, embed_dim=1536, patch=14)
  Multi-scale features extracted at layers 5 / 11 / 17 / 23
  Projected to 256 / 512 / 1024 / 2048 channels via Conv2d
       │
       ├──► UperNet decoder → semantic logits (B, 6, H, W)
       │
       └──► UperNet decoder → HV maps       (B, 2, H, W)
```

**UNI2 backbone** — `vit_giant_patch14_224` (timm), loaded from `MahmoodLab/UNI2-h` pretrained weights, 1.1 B parameters.

**UNI2UperHoVer** — dual-head UperNet decoder; semantic head classifies tissue type per pixel; HV head produces horizontal/vertical nuclear centroid distance fields used for marker-controlled watershed instance segmentation.

---

## Installation

```bash
# Core (model architecture only)
pip install segtme-uni2

# With large-tile inference pipeline (LargeTilePredictor)
pip install "segtme-uni2[predict]"
```

Core dependencies: `torch`, `transformers`, `timm`, `safetensors`, `huggingface_hub`.  
Predict extras: `numpy`, `pillow`, `torchvision`, `opencv-python`, `scipy`, `scikit-image`.

---

## Usage

```python
from segtme import UNI2UperHoVer

# Load from HuggingFace (downloads weights automatically)
model = UNI2UperHoVer.from_pretrained("mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-0")
model.eval().cuda()

# Forward pass — input: (B, 3, H, W), ImageNet normalised, recommended 224×224 tiles
import torch
pixel_values = torch.randn(1, 3, 224, 224).cuda()
sem_logits, hv_maps = model(pixel_values)

# sem_logits: (B, 6, H, W)  — per-class tissue logits
# hv_maps:    (B, 2, H, W)  — horizontal / vertical distance gradients
```

### Output classes

| Channel | Class |
|---|---|
| 0 | Background |
| 1 | Neoplastic |
| 2 | Inflammatory |
| 3 | Connective |
| 4 | Dead |
| 5 | Epithelial |

### Recommended inference scale

Resize input tiles so one pixel corresponds to the target MPP before running inference:

| Model | Target MPP | Downscale factor |
|---|---|---|
| M1 — PanNuke | 0.25 µm/px | `image_mpp / 0.25` |
| M2 — TCGA-UT-0 | 0.314 µm/px | `image_mpp / 0.314` |
| M3 — TCGA-UT-012345 | 0.5 µm/px | `image_mpp / 0.5` |

### Large-tile inference with LargeTilePredictor

`LargeTilePredictor` handles the full end-to-end pipeline on images of arbitrary size:
upscaling, tiling (224×224, 50% overlap), stitching, watershed instance segmentation,
and downscaling back to original resolution.

```python
# Requires: pip install "segtme-uni2[predict]"
from segtme import UNI2UperHoVer, LargeTilePredictor

# Load from HuggingFace (downloads weights automatically)
model = UNI2UperHoVer.from_pretrained("mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-0")

predictor = LargeTilePredictor(
    model,
    input_mpp=0.5,    # µm/px of input image (0.5 = 20× scanner)
    model_mpp=0.314,  # M2 training resolution (default)
)

result = predictor.predict("path/to/image.png")

result.sem        # uint8  (H, W)    6-class semantic mask
result.inst       # uint16 (H, W)    watershed instance map (0 = background)
result.class_img  # uint8  (H, W, 3) RGB instances coloured by class
result.outline    # uint8  (H, W, 3) original H&E + class-coloured boundaries

# Save outputs
from PIL import Image
import cv2
Image.fromarray(result.sem).save("pred_sem.png")
Image.fromarray(result.inst).save("pred_inst.png")
Image.fromarray(result.outline).save("pred_outline.png")
```

Load from a local safetensors checkpoint (no HuggingFace download):

```python
from segtme import UNI2UperHoVer, UNI2UperHoVerConfig, LargeTilePredictor
from safetensors.torch import load_file

config = UNI2UperHoVerConfig()
model  = UNI2UperHoVer(config)
model.load_state_dict(
    load_file("checkpoint-212500/model.safetensors"), strict=True)

predictor = LargeTilePredictor(model, input_mpp=0.5)
result = predictor.predict("image.png")
```

Model MPP reference for `model_mpp`:

| Model | `model_mpp` | Best for |
|---|---|---|
| M1 — PanNuke | `0.25` | 40× / ~0.25 µm/px inputs |
| M2 — TCGA-UT-0 | `0.314` *(default)* | 20× / 0.5 µm/px inputs |
| M3 — TCGA-UT-012345 | `0.5` | Variable 0.5–1.0 µm/px inputs |

### Low-level tiling helpers

The standalone functions underlying `LargeTilePredictor` are also exported for
users who need to compose the pipeline manually:

```python
from segtme import (
    pad_to_multiple, tile_image, stitch_semantic, stitch_hv,
    predict_tile_dual, clean_mask, run_watershed,
    map_instances_to_classes, build_outline_image,
    REGION_COLORS, REGION_LABELS,
)
```

---

## Training Curriculum

Three-stage curriculum training, each stage initialised fresh (no weight inheritance):

| Stage | Model | Dataset | Epochs | Steps | mIoU |
|---|---|---|---|---|---|
| 1 | M1 | PanNuke | 249 | 24,651 | 0.9313 |
| 2 | M2 | TCGA-UT subset 0 | 250 | 212,500 | 0.8197 |
| 3 | M3 | TCGA-UT subsets 0–5 | 100 | 335,100 | 0.7724 |

All stages: initial LR 5×10⁻⁵, linear decay, AdamW optimiser. Backbone: UNI2-h (frozen or fine-tuned depending on stage).

---

## Citation

If you use this model in your work, please cite:

```bibtex
@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}
}
```

---

## Links

- Paper (arXiv): https://arxiv.org/abs/2606.17702
- PyPI: https://pypi.org/project/segtme-uni2/
- HuggingFace (M1): https://huggingface.co/mizjaggy18/SegTME-UNI2-UperHoVer_PanNuke
- HuggingFace (M2): https://huggingface.co/mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-0
- HuggingFace (M3): https://huggingface.co/mizjaggy18/SegTME-UNI2-UperHoVer_TCGA-UT-012345
- UNI2 backbone: https://huggingface.co/MahmoodLab/UNI2-h
