ClementP/ChoroidSegmentationModule

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):

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 (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