aungmthein/oct-cell-layer-fold4-swin-umamba

OCT cartilage cell-layer (fold 4, Swin-UMamba)

Production checkpoint for the OCT_Fabiola cell-layer segmentation task. 3-class semantic segmentation (above / cell layer / below) on 2D OCT images of engineered cartilage constructs.

Architecture

• Encoder: Swin-UMamba VMamba-Tiny, ImageNet-pretrained (Liu et al., 2024)
• Head: 3-class softmax + per-pixel CE
• Trained: 100 epochs, 384x384 random crops, fold 4 of 5-fold CV with construct-grouped splits

Provenance

• Run id: phase4-1f3f21ce
• Trained 2026-05-21 on NVIDIA GB10 (DGX Spark)

Numbers (5-fold CV + cross-batch OOD)

slice	n	Dice	MAE (px)	ECE	Brier
val (fold 4)	14	0.84	32	0.05	0.13
OOD aggregate	44	0.572	104.47	0.174	0.432
OOD restricted (thin-bounded subset)	29	~0.81	~48	--	--
OOD no-bottom subset	15	~0.10	n/a	--	--

Important caveat on the OOD numbers. The OOD slice is bimodal-by-regime: 15 of 44 OOD images are in the no-bottom regime (cell layer extends below the visible frame), which is completely absent from training (0/68 training images). The model has zero training signal for that regime, so the aggregate OOD Dice 0.572 is dominated by an extrapolation failure, not a generalization failure on the in-distribution regime. The honest metric is the restricted-OOD number on the thin-bounded subset.

Usage

from huggingface_hub import hf_hub_download

ckpt_path = hf_hub_download(
    repo_id="aungmthein/oct-cell-layer-fold4-swin-umamba",
    filename="phase4-1f3f21ce_softmax_B__swin_umamba_fold4.pth",
)
# Loader integration: see the demo Space at
# https://huggingface.co/spaces/aungmthein/oct-cartilage-cell-layer

Inputs

• 2D OCT image, single-channel grayscale recommended.
• Cross-scale: the loader handles any H x W (zero-padded internally to a multiple of 32).
• Recommended physical pixel spacing: 2.0 um/px (baseline 8000 x 1500 / 16 x 3 mm).

Outputs

• 3-class mask with values {0: above, 1: cell layer, 2: below}.
• Thickness statistics (mean +/- std, percentiles, columns-measured count).

License

MIT (matches the upstream Swin-UMamba reference).