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