dinoct

DINO-style self-supervised pretraining for OCT B-scan images, plus curve head (LoRA) post-train stage.

Repo layout

• dinoct/: Python package (models, data, training)
• configs/: YAML configs (merged: configs/ssl_default_config.yaml + configs/train/oct.yaml)
• eval/: evaluation entrypoints
• tools/: export and dataset-maintenance utilities

Quick Start

• Python: >=3.12
• CUDA: single GPU

This repo uses uv:

uv sync
uv run python -m dinoct --help

Data

Default expected layout under data/oct/:

• data/oct/raw/*.jpg
• data/oct/background/*.jpg
• data/oct/labeled/<image_stem>.txt (optional; marks an image as labeled)
• data/oct/extra/entries.npy (metadata cache; regenerated each run)
• data/oct/extra/splits.csv (paper train/val/test assignment)

Each label file should contain either:

• 500 floats (one per column), or
• a 500×2 table (x, y) (the second column is used).

You can change the dataset paths via the dataset string: OCT:root=<root>[:extra=<extra>] (see configs/train/oct.yaml). In that string, root/extra refer to dataset directories (not the repo root). The dataset name token is case-insensitive.

Splits

Replication experiments expect data/oct/extra/splits.csv. SSL pretraining uses split=train, while post-train reads the explicit train and val assignments from that file.

If splits.csv is already present, you can skip split generation. If you need to regenerate it, run:

uv run python tools/data/build_oct_manifest.py --dir data/oct
uv run python tools/data/build_oct_splits.py --dir data/oct --seed 0 --train-frac 0.7 --val-frac 0.15

Replication

Use the single replication entrypoint:

uv run python eval/run_paper.py --help

All paper results use the checked-in default config at configs/train/oct.yaml. The replication workflow below is intended to be run without extra optimizer or schedule overrides.

Recommended replication order:

1. Make sure data/oct/extra/splits.csv exists. If it is already present, skip this step.

uv run python tools/data/build_oct_manifest.py --dir data/oct
uv run python tools/data/build_oct_splits.py --dir data/oct --seed 0 --train-frac 0.7 --val-frac 0.15

2. Make sure the default pretrained backbone exists at:

outputs/pretrain/dinov3_pretrain.pth

You can grab it from Github releases or generate it by running the default training config once:

uv run python -m dinoct \
  --config configs/train/oct.yaml \
  --output-dir outputs

This writes both outputs/pretrain/dinov3_pretrain.pth and the default outputs/post_train/ checkpoint. eval/run_paper.py reuses the backbone checkpoint from outputs/pretrain/dinov3_pretrain.pth.

3. Run the paper tables and benchmarks from eval/run_paper.py.

Main paper commands:

uv run python eval/run_paper.py low-data --resume
uv run python eval/run_paper.py ablations --resume
uv run python eval/run_paper.py robustness
uv run python eval/run_paper.py skin-db-export --overwrite
uv run python eval/run_paper.py skin-db-paper-dl
uv run python eval/run_paper.py skin-db-cv --layer upper --resume
uv run python eval/run_paper.py skin-db-cv --layer dej --resume

What each command produces:

• low-data
  • runs the DINOCT / UNET / FCBR data-efficiency experiments
  • writes outputs under outputs/data_efficiency/
  • fills the paper tables under outputs/paper_tables/data_efficiency_*
• ablations
  • runs the final single-component ablation around configs/train/oct.yaml
  • writes outputs/ablations/final_component_ablation/final_component_ablation.csv
  • writes outputs/ablations/final_component_ablation/final_component_ablation.tex
• robustness
  • runs clean + synthetic artifact evaluation for the default DINOCT checkpoint
  • writes under outputs/robustness/dinoct/test/
• skin-db-export
  • exports skin_db.mat into:
    • data/skin_db/upper_gt
    • data/skin_db/dej_gt
• skin-db-paper-dl
  • evaluates the paper-provided dl predictions from skin_db.mat
  • writes outputs/skin_db_mat_dl_eval/paper_dl_summary.csv
• skin-db-cv --layer upper|dej
  • runs the patient-wise external CV comparison for DINOCT and the paper-style U-Net baseline
  • writes under:
    • outputs/frontiers_cv_upper_gt/
    • outputs/frontiers_cv_dej_gt/

Expected external-data input:

• data/skin_db/skin_db.mat
• data/skin_db/plot_results.m
• data/skin_db/readme.txt

Evaluation

Evaluate a fused curve checkpoint on val or test using the paper metrics:

• MAE
• P95
• signed Bias
• |Bias|
• Acc@2px
• Acc@4px
• SpikeRate@kappa

UV_CACHE_DIR=/tmp/uv-cache uv run python eval/evaluate_curve.py \
  --config configs/train/oct.yaml \
  --curve-ckpt outputs/post_train/fused_curve_best.pth \
  --split test

This writes:

• per_scan_metrics.csv
• per_recording_metrics.csv
• summary.json

If --spike-kappa is not provided, the evaluator estimates kappa from the reference centerlines on the validation split using the configured quantile.

Learned Baselines

Train a supervised learned baseline with the same OCT splits, Gaussian targets, and evaluation metrics:

UV_CACHE_DIR=/tmp/uv-cache uv run python eval/train_learned_baseline.py \
  --config configs/train/oct.yaml \
  --model-type unet

or

UV_CACHE_DIR=/tmp/uv-cache uv run python eval/train_learned_baseline.py \
  --config configs/train/oct.yaml \
  --model-type fcbr

Both scripts write checkpoints under outputs/learned_baselines/<model_type>/. Evaluate them with the same curve evaluator:

UV_CACHE_DIR=/tmp/uv-cache uv run python eval/evaluate_curve.py \
  --curve-ckpt outputs/learned_baselines/unet/curve_best.pth \
  --split test

Labeling (curve editor)

The interactive curve label editor requires matplotlib:

uv sync --extra label
uv run python tools/data/curve_labeler.py --dir data/oct

Direct Training

Train the default OCT recipe end to end:

uv run python -m dinoct \
  --config configs/train/oct.yaml \
  --output-dir outputs/run1

This uses the values already defined in configs/train/oct.yaml and runs SSL pretrain followed by post-train.

Outputs:

• outputs/run1/pretrain/dinov3_pretrain.pth
• outputs/run1/pretrain/train.log, metrics.csv, config_used.yaml
• outputs/run1/post_train/fused_curve.pth
• outputs/run1/post_train/fused_curve_best.pth
• outputs/run1/post_train/val_summary.json

To rerun only the curve post-train stage with an existing backbone, keep the same config and use:

uv run python -m dinoct \
  --config configs/train/oct.yaml \
  --output-dir outputs/run1 \
  --post-train-only \
  --pretrained-backbone outputs/run1/pretrain/dinov3_pretrain.pth

Export (TorchScript/ONNX)

uv run python tools/export_model.py --model outputs/run1/post_train/fused_curve_best.pth --outdir exports

Dataset

Dataset are available from https://huggingface.co/datasets/rjbaw/oct and using the huggingface dataset loader.
Set train.dataset_path to OCT:hub=rjbaw/oct (or include root=... for local-first).

License

Apache-2.0.

This project includes alot of code derived from Meta Platforms, Inc. and affiliates' DINOv2 and DINOv3 repositories, licensed under the Apache License, Version 2.0.