README.md

---
license: cc-by-nc-4.0
pipeline_tag: image-segmentation
tags:
  - sam2
  - segment-anything
  - medical-imaging
  - optical-coherence-tomography
  - oct
  - glaucoma
  - image-segmentation
  - pytorch
base_model: facebook/sam2.1-hiera-base-plus
---

# SAM2-OCT — Fine-tuned SAM2 checkpoints for OCT segmentation

Fine-tuned [SAM 2 (Segment Anything Model 2)](https://github.com/facebookresearch/sam2)
checkpoints for interactive, multi-class segmentation of retinal **Optical Coherence
Tomography (OCT)** images. Developed as part of an MSc dissertation at the University
of the Witwatersrand.

These weights are designed to be used with the companion annotation tool
**[CVAT-OCT](https://github.com/enslinr/cvat-oct)** (a fork of CVAT with a SAM2-OCT
serverless interactor).

> ⚠️ **Research use only.** These models are experimental research artifacts and are
> **not** a medical device. They must **not** be used for clinical diagnosis, screening,
> or treatment decisions.

## Model description

These checkpoints adapt SAM 2 for **multi-class retinal layer segmentation** using a
*semantically aware* modification: SAM's generic mask tokens are replaced with
**dedicated per-layer mask tokens** and per-class output heads, so that every retinal
layer class is predicted in a **single forward pass**, while SAM 2's interactive
prompting interface is preserved for optional manual refinement. The image encoder is
SAM 2.1's Hiera Base+ backbone, fine-tuned end-to-end together with the modified mask
decoder.

Key properties:

- **Single-pass multi-class output** — one mask channel per retinal layer, rather than
  one binary mask per prompt.
- **Interactive-ready** — point / box / rough-mask prompting is retained for
  human-in-the-loop correction (see the `MGU_prompted` checkpoint).
- **Data-efficient** — on the macular (NR206) task the approach substantially
  outperforms a purpose-built specialised baseline when annotated data is scarce.

## Checkpoints

| File | Description | Base |
|------|-------------|------|
| `MGU/final_runs_Glaucoma_last.pt` | Semantically aware SAM2 trained on the **MGU** peripapillary (glaucoma) dataset. Automatic single-pass segmentation of nine retinal layers plus the optic-disc region (ten foreground classes + background) on peripapillary OCT B-scans. | SAM2.1 Hiera Base+ |
| `MGU_prompted/MGU_prompt_training_last.pt` | **Prompted** variant of the MGU model. Adds class-aware point and rough-mask prompt encoders so a reviewer can interactively guide or correct the output. Its automatic prediction matches the standard MGU model; brushing a rough mask over an error region improves the local segmentation (≈ +6.5% mIoU in the automatic-prompt evaluation), whereas point prompts did not yield a measurable gain in the current form. | SAM2.1 Hiera Base+ |
| `NR206/final_runs_NR206_last.pt` | Semantically aware SAM2 trained on the **NR206** macular dataset (healthy eyes). Automatic single-pass segmentation of eight retinal layers (+ background) on macular OCT B-scans. | SAM2.1 Hiera Base+ |

Each checkpoint is ~880 MB.

> The base `sam2.1_hiera_base_plus.pt` checkpoint is **not** included here — download it
> from Meta's [SAM2 releases](https://github.com/facebookresearch/sam2/releases). Only the
> fine-tuned OCT weights are hosted in this repository.

## Intended use

- Interactive / automatic segmentation of OCT structures within the CVAT-OCT tool.
- Research and educational exploration of SAM2 for medical image segmentation.

### Out of scope

- Any clinical, diagnostic, or patient-facing use.
- Deployment on imaging modalities or populations other than those it was trained on
  (results are not expected to transfer). In particular, each checkpoint is specialised
  to its dataset's acquisition device, scan region (macular vs. peripapillary), and label
  set; cross-device / cross-region generalisation is not guaranteed.

## How to use (with CVAT-OCT)

Download the checkpoint(s) into the matching `models/` sub-directory of the
SAM2-OCT serverless function, then start CVAT-OCT:

```bash
# from the root of a cvat-oct clone
mkdir -p serverless/pytorch/sam2-OCT-interactor/models/MGU

# Option A: huggingface_hub (recommended)
pip install huggingface_hub
python - <<'PY'
from huggingface_hub import hf_hub_download
hf_hub_download(
    repo_id="enslinr/sam2-oct",                     # <-- your HF repo id
    filename="MGU/final_runs_Glaucoma_last.pt",
    local_dir="serverless/pytorch/sam2-OCT-interactor/models",
)
PY

# Option B: direct download
# wget https://huggingface.co/enslinr/sam2-oct/resolve/main/MGU/final_runs_Glaucoma_last.pt \
#   -O serverless/pytorch/sam2-OCT-interactor/models/MGU/final_runs_Glaucoma_last.pt
```

Point the function at the checkpoint via the `SAM2_CHECKPOINT` environment variable
(see `serverless/pytorch/sam2-OCT-interactor/function.yaml` and `docker-compose.override.yml`
in the CVAT-OCT repo). A short video walkthrough of the end-to-end annotation workflow is
available at <https://enslinr.github.io/cvat-oct/>.

## Training data

All training data are **publicly available, fully anonymised** OCT datasets; no new human
data were collected. Images are grayscale OCT B-scans.

**NR206 (macular, healthy eyes)** — 206 macular B-scans of healthy human eyes, derived from
the OCTID database. Acquired with a Cirrus HD-OCT device (Carl Zeiss Meditec) using an
840 nm source (≈ 5 µm axial resolution); original resolution 500 × 750 px. Labels cover
**8 retinal-layer classes** (NFL, GCL+IPL, INL, OPL, ONL, ELM+IS, OS, RPE) plus background.
Author-provided split: 126 train / 40 val / 40 test.
Dataset: He et al., *Frontiers in Bioengineering and Biotechnology*, 2023 (NR206).

**MGU (peripapillary, glaucoma)** — peripapillary OCT from 61 subjects (Shanghai General
Hospital), acquired with a DRI OCT-1 Atlantis device (Topcon) over a 20.48 × 7.94 mm field
centred on the optic nerve head; original resolution 1024 × 992 px. 122 manually annotated
B-scans covering **10 foreground classes** — nine retinal layers (RNFL, GCL, IPL, INL, OPL,
ONL, IS/OS, RPE, Choroid) and the optic-disc region — plus background. Author-provided
split: 74 train / 24 val / 12 test. Dataset: Li et al., *Biomedical Optics Express*, 2021 (MGU).

**Preprocessing.** B-scans are resized to SAM 2's 1024 × 1024 input (stretch resizing,
selected from a resize-strategy comparison), the single grayscale channel is duplicated
across the three input channels, and images are normalised with ImageNet mean/standard
deviation. Training used data augmentation (rotation, brightness/contrast jitter, Gaussian
noise and blur, elastic and grid distortion, gamma adjustment, and CLAHE).

**Ethics / data use.** The study used only publicly available, de-identified datasets and
collected no new human data; the University of the Witwatersrand granted a waiver of ethics
clearance (**Ethics Waiver Number: WCSAM-2024-19**). Because the models are derived solely
from these public datasets, releasing the fine-tuned weights is consistent with that use.
Users should nonetheless comply with the terms of the underlying NR206 and MGU datasets.

## Training procedure

- Fine-tuned end-to-end from **SAM 2.1 Hiera Base+**, with SAM's mask tokens replaced by
  per-layer tokens and per-class MLP output heads.
- **Loss:** a combined objective (Focal + Soft Dice + Soft IoU).
- **Optimisation:** separate learning rates for the mask decoder (≈ 7.2 × 10⁻³) and the
  image encoder (≈ 2.2 × 10⁻⁷), AdamW-style weight decay 0.01, gradient clipping 2.0,
  cosine schedule with warmup. Final hyperparameters were selected via a 110-run Bayesian
  sweep (Weights & Biases) optimising validation mIoU.
- **Hardware:** a single NVIDIA GeForce RTX 3090.
- The `MGU_prompted` checkpoint additionally trains class-aware sparse (point) and dense
  (rough-mask) prompt encoders so that interactive prompts can target a specific layer.

## Evaluation

Models are evaluated on the authors' **held-out test splits** using per-layer Dice, mean
IoU (mIoU), and mean Dice. Statistical comparisons against retrained specialised baselines
(EMV-Net and LightReSeg) use two-sided Wilcoxon signed-rank tests on per-image mIoU.

**NR206 test set (macular, healthy; n = 40).** Our model attains the highest score on every
aggregate and per-layer metric, significantly outperforming both retrained baselines
(mIoU and mean Dice, *p* < 0.001).

| Metric | mIoU | Dice | NFL | GCL+IPL | INL | OPL | ONL | ELM+IS | OS | RPE |
|---|---|---|---|---|---|---|---|---|---|---|
| Ours (SAM2) | **85.6** | **92.0** | 91.5 | 96.6 | 91.4 | 83.5 | 95.5 | 92.7 | 88.4 | 96.4 |

*(mIoU/Dice are the aggregate scores; the remaining columns are per-layer Dice.)*

**MGU test set (peripapillary, glaucoma; n = 48).** Our model significantly outperforms the
retrained baselines on aggregate mIoU (*p* < 0.001) and matches the purpose-built published
EMV-Net to within 0.2 mIoU, despite being a general-purpose foundation-model adaptation.

| Metric | mIoU | Dice | RNFL | GCL | IPL | INL | OPL | ONL | IS/OS | RPE | Choroid | Disc |
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Ours (SAM2) | 68.6 | 80.4 | 81.6 | 65.9 | 70.8 | 76.0 | 80.3 | 90.7 | 85.8 | 81.7 | 89.3 | 82.2 |

The same approach was additionally evaluated on a diabetic macular oedema dataset and on
combined multi-dataset training (approaching a purpose-built universal baseline); those
results are reported in the dissertation but the corresponding checkpoints are not released
here. See the dissertation for full tables, per-image statistics, ablations, and the
prompted-refinement study.

## Limitations

- **Single-run point estimates** for some development comparisons; final test-set numbers
  above are single-run results interpreted against measured seed-to-seed variability
  (≈ 0.1–0.3 mIoU).
- **Dataset-specific.** Each checkpoint is trained and evaluated on one dataset/device;
  performance on other devices, protocols, or pathologies is not expected to transfer.
- **Input resolution.** SAM 2's fixed 1024 × 1024 input requires upscaling OCT B-scans,
  which can introduce interpolation artefacts affecting fine boundary precision.
- **No pathology segmentation** beyond the labelled layer/disc classes (e.g. drusen or
  fluid are not segmented by the released macular/glaucoma checkpoints).

## Base model & license

- Fine-tuned from **SAM 2.1 Hiera Base+** (`facebook/sam2.1-hiera-base-plus`), released by
  Meta AI under the Apache-2.0 license.
- These fine-tuned weights are released under **CC BY-NC 4.0** (attribution, non-commercial).
  Use of the weights must also respect the terms of the underlying public NR206 and MGU
  datasets.

## Citation

If you use these checkpoints, please cite the dissertation and this repository, along with
the CVAT-OCT project (see its `CITATION.cff`) and the upstream **SAM 2** and **CVAT** projects.

```bibtex
@mastersthesis{roux_sam_oct_2026,
  author  = {Roux, Enslin},
  title   = {Adapting the Segment Anything Model (SAM) for Retinal OCT Layer Segmentation},
  school  = {University of the Witwatersrand},
  year    = {2026}
}

@software{roux_cvat_oct,
  author  = {Roux, Enslin},
  title   = {CVAT-OCT: AI-assisted segmentation of OCT images (a CVAT fork)},
  year    = {2026},
  url     = {https://github.com/enslinr/cvat-oct}
}
```

## Author

Enslin Roux — University of the Witwatersrand.