---
base_model: timm/convnext_tiny.in12k_ft_in1k
library_name: pytorch
pipeline_tag: image-segmentation
tags:
- medical-imaging
- ophthalmology
- oct
- choroid
- segmentation
- upernet
- convnext
- segmentation-models-pytorch
- pytorch_model_hub_mixin
- model_hub_mixin
---

# Choroid segmentation in OCT B-scans (version 3.0.0)

A binary segmentation model for the **choroid** in OCT B-scans: it labels the
region between Bruch's membrane and the choroid-sclera interface (CSI). It is the
first stage of the OcularRigidity pipeline (choroid segmentation, then video
registration, then pulsation analysis). Its frozen encoder is also the feature
extractor of the registration model `ClementP/OCTVideoRegistration`.

| | |
|---|---|
| Architecture | UPerNet decoder, ConvNeXt-Tiny encoder (`tu-convnext_tiny`, ImageNet init), from `segmentation_models_pytorch` |
| Parameters | 37.0 M |
| Input | 1 channel (grayscale B-scan) at the **PLEX Elite pixel size: 1.95 µm axial × 5.9 µm lateral** |
| Output | 1 channel of logits; choroid = `logit > 0` |
| Revision | `version-3.0.0` |

## Usage

The weights load with the `ChoroidSegmentationModule` class from the
OcularRigidity package (`ocularrigidity.segmentation.trainer.pl_module`):

```python
import torch
from ocularrigidity.segmentation.trainer.pl_module import ChoroidSegmentationModule

model = ChoroidSegmentationModule.from_pretrained(
    "ClementP/ChoroidSegmentationModule", revision="version-3.0.0"
).eval().cuda()

# bscan: uint8 array (H, W), already at 1.95 µm/px axially and 5.9 µm/px laterally,
# H and W padded to multiples of 32
x = torch.from_numpy(bscan).float().div(255).sub(0.5).div(0.5)[None, None].cuda()
with torch.no_grad():
    mask = model(x)[0, 0] > 0  # (H, W) bool
```

**Preprocessing**
- **Pixel size:** resample images from other devices to the PLEX pixel size first.
  Training used only small scale jitter (±10%), so the model expects that scale.
- **Normalisation:** `(x / 255 - 0.5) / 0.5`.


## Training data

All images were resampled to the PLEX pixel size. The splits are made at the patient level.

| Source | Device | Train | Val | Test |
|---|---|---|---|---|
| In-house, annotated (Didascalie) | Zeiss PLEX Elite 9000 | 1,769 | 118 | 329 |
| In-house, annotated (Didascalie) | Heidelberg Spectralis (videos) | 327 | 24 | 77 |
| In-house, annotated (Didascalie) | Heidelberg Spectralis | 81 | 10 | 18 |
| [OIMHS](https://www.nature.com/articles/s41597-023-02675-1) (choroid label) | Spectralis | 125 patients, 3,859 B-scans in total, split 0.8 / 0.05 / 0.15 by patient (seed 42) | | |

The in-house annotations are not public.

## Training procedure

**Setup**
- **Loss:** Dice + cross-entropy, plus a thickness-smoothness term (weight 0.1)
  and a gap-detection term (weight 0.05) that penalise jagged or broken choroid masks.
- **Optimiser:** AdamW, learning rate 1e-4, weight decay 1e-5. Cosine schedule
  with 2 warm-up epochs, 60 epochs, no early stopping.
- **Batches:** batch size 8, random 1024 × 768 crops, bf16 mixed precision, one GPU.

**Model selection:** the architecture was chosen in a W&B sweep over SegFormer,
UPerNet, U-Net, U-Net++ and FPN variants trained with the same recipe. UPerNet
ConvNeXt-T gave the best trade-off between Dice and boundary errors.

**Checkpoint:** this release is the **final-epoch** checkpoint, not the best-Dice
one. The final checkpoint has a clearly lower thickness error for almost the same Dice.

## Evaluation

Validation split (in-house plus OIMHS), final checkpoint:

| Metric | Value |
|---|---|
| Dice | 0.958 |
| Choroid thickness MAE (per A-scan) | 12.8 µm |
| CSI position MAE (per A-scan) | 19.5 µm |

**How the errors are computed**
- Errors are per column, over the columns where the ground truth has choroid, at 1.95 µm/px.
- The CSI is the deepest choroid pixel of each column.

## Limitations

- **Single frames:** the model segments each B-scan independently, with no
  temporal consistency across video frames. In the pipeline, consistency comes
  from the registration stage.
- **Training domain:** trained on PLEX Elite and Spectralis only, with mostly
  macular and some optic nerve head scans. Other devices, or scans with
  pathologies that change choroid appearance, may need validation.
- **CSI:** the choroid-sclera interface is the hardest boundary, and its error is
  larger than the thickness error.
- **Research use only.** This is not a medical device.

## Version history

- `version-3.0.0`:
  - UPerNet ConvNeXt-T;
  - new annotated set plus OIMHS, everything at the PLEX pixel size;
  - final-epoch checkpoint.
- `version-2.0.0`: previous model, kept on its branch.

## Citation

If you use this model, please cite the OcularRigidity work (reference to come)
and the OIMHS dataset:

```
Ye, X., He, S., Zhong, X. et al. OIMHS: An Optical Coherence Tomography Image Dataset Based on Macular Hole Manual Segmentation. Sci Data 10, 769 (2023). https://doi.org/10.1038/s41597-023-02675-1
```