Uncertainty-aware Retinal Layer Segmentation in OCT through Probabilistic Signed Distance Functions

Paper  |  MIDL 2024 / PMLR  |  Project page  |  Code

This repository contains the official implementation of uncertainty-aware retinal layer segmentation in optical coherence tomography (OCT) using probabilistic signed distance functions (pSDFs).

Figure 1 from the paper: signed distance functions represent retinal layer geometry through level sets, while probabilistic modeling provides spatially meaningful uncertainty.

Highlights

• Geometric segmentation: retinal layer boundaries are represented by signed distance functions and their level sets.
• Probabilistic predictions: the model estimates a mean and variance for the signed distance representation.
• Uncertainty-aware analysis: uncertainty can reveal ambiguous, noisy, or pathological regions in OCT scans.
• Robustness experiments: the project includes settings for evaluating synthetic artifacts and noise.
• Hydra configuration: models, datasets, losses, experiments, and logging are configured with composable YAML files.

Repository structure

.
├── config/          # Hydra configuration files and experiment settings
├── dataloader/      # DataLoader construction and sampling utilities
├── dataset/         # Dataset definitions and data-related code
├── model/           # Network architectures and model handlers
├── utils/           # Losses, logging, writers, and helper functions
├── trainer.py       # Training entry point
├── inference.py     # Checkpoint-based inference entry point
├── test.py          # Data-loading and evaluation utilities
├── environment.yml  # Conda environment specification
└── requirements.txt # Python package requirements

Installation

The provided environment targets Linux, Python 3.8, PyTorch 1.13, and CUDA-enabled execution. A GPU is recommended for training.

Conda

git clone https://github.com/niazoys/RLS_PSDF.git
cd RLS_PSDF

conda env create -f environment.yml
conda activate pytorch-template

Alternatively, install the package dependencies with pip:

pip install -r requirements.txt

If you install PyTorch separately, choose a version and CUDA runtime compatible with your system. The pinned environment specifies torch==1.13.0 and torchvision==0.14.0.

Docker

A CUDA 11.3-based Dockerfile is included for users who prefer a containerized setup:

docker build -f DockerFile -t rls-psdf .
docker run --gpus all -it --rm -v "$PWD":/app rls-psdf

Check the Dockerfile and host NVIDIA Container Toolkit configuration before launching GPU workloads.

Data and configuration

The OCT datasets are not included in this repository. Prepare the data separately and update the Hydra configuration with the appropriate data locations and train/validation/test splits.

The main configuration is config/default.yaml. Dataset, model, and experiment configurations are organized under the corresponding subdirectories of config/.

At minimum, set the data root to your prepared dataset:

data:
  data_root_dir: /path/to/your/data

Before training or inference, also verify the dataset-specific configuration, the number of output classes/layers, checkpoint paths, and the selected device.

Training

Run the default experiment with:

python trainer.py

Hydra experiment configurations can be selected with a positional argument:

python trainer.py experiment=<experiment_name>

For example, for config/experiment/example.yaml:

python trainer.py experiment=example

Configuration values can be overridden from the command line:

python trainer.py \
  device=cuda \
  data.data_root_dir=/path/to/your/data \
  train.num_epoch=100

The trainer creates Hydra run directories under outputs/ and can log experiments through Weights & Biases when enabled in the configuration:

wandb login

Inference

Run inference with the same Hydra configuration system:

python inference.py experiment=<experiment_name>

Before running inference, make sure that:

• the test data directory is configured;
• the model architecture matches the checkpoint;
• load.resume_state_path or load.network_chkpt_path points to the desired checkpoint; and
• inference mode and output-saving options are enabled in the selected configuration.

The inference entry point supports CPU, single-GPU, and distributed GPU execution according to the active configuration.

Project page

The project website includes the paper figures, an intuitive explanation of the SDF and uncertainty formulation, qualitative examples, comparisons, and the complete BibTeX entry:

https://niazoys.github.io/RLS_PSDF/

Citation

If you use this code or the probabilistic signed distance function approach, please cite:

@InProceedings{pmlr-v250-islam24a,
  title     = {Uncertainty-aware retinal layer segmentation in OCT through probabilistic signed distance functions},
  author    = {Islam, Mohammad Mohaiminul and de Vente, Coen and Liefers, Bart and Klaver, Caroline and Bekkers, Erik J. and S{\'a}nchez, Clara I.},
  booktitle = {Proceedings of The 7th International Conference on Medical Imaging with Deep Learning},
  pages     = {672--693},
  year      = {2024},
  volume    = {250},
  series    = {Proceedings of Machine Learning Research},
  publisher = {PMLR},
  url       = {https://proceedings.mlr.press/v250/islam24a.html}
}

License

No license file is currently included in this repository. Please contact the authors before redistributing the code or using it in a commercial product.