vit-oct-wamd

A SigLIP vision transformer fine-tuned to classify optical coherence tomography (OCT) B-scan images as normal or wet AMD (neovascular age-related macular degeneration).

Released as supporting material for an accompanying paper.

Model description

• Base encoder: google/siglip-so400m-patch14-384, fine-tuned end-to-end.
• Head: a single linear layer (1152 → 2) on top of the encoder's CLS token, with dropout before the head.
• Auxiliary branch: the released checkpoint also includes a frozen T5-base text encoder + projection (siglip_loss in the state dict), used only for an auxiliary contrastive image-text alignment loss during training. It is not used at inference time — forward() only calls image_encoder and cls_head.

Dataset

Dataset: kavin-aravindhan/oct-wamd. Public access is pending IRB approval — for interim access, contact Kavin Aravindhan (kr3131@columbia.edu).

Files in this repo

File	Purpose
best_model.pt	Full training checkpoint (classifier weights + auxiliary branch + optimizer/scheduler state).
modeling.py, alignment.py, embedder.py	Model definition matching the checkpoint exactly.
infer.py	Quick-start CLI: run the classifier on one or more images.
demo_app.py	Local interactive Gradio demo (image upload → prediction).
train.py	Trainer script reproducing the paper's training procedure.
training_config.json	Exact hyperparameters used for this checkpoint.
requirements.txt, requirements-train.txt, requirements-demo.txt	Pinned dependency versions for infer.py, train.py, and demo_app.py respectively.

Quick start

pip install huggingface_hub
huggingface-cli download kavin-aravindhan/vit-oct-wamd infer.py requirements.txt --local-dir .
pip install -r requirements.txt
python infer.py path/to/scan.png

requirements.txt pins torch/transformers/huggingface_hub to the exact versions used to produce best_model.pt — an unpinned pip install transformers can resolve to a version whose SiglipVisionModel uses different internal parameter names, which will fail to load this checkpoint's state dict.

path/to/scan.png: wet_amd (confidence=0.513)

First run downloads the checkpoint (~5.1GB) and caches it; subsequent runs take a few seconds per image on GPU, ~15-20s on CPU.

Or from Python:

from huggingface_hub import hf_hub_download
import sys, os, torch

repo_id = "kavin-aravindhan/vit-oct-wamd"
local_dir = os.path.dirname(hf_hub_download(repo_id=repo_id, filename="modeling.py"))
for f in ["alignment.py", "embedder.py", "best_model.pt"]:
    hf_hub_download(repo_id=repo_id, filename=f, local_dir=local_dir)

sys.path.insert(0, local_dir)
from modeling import load_model, IMAGE_SIZE

model = load_model(os.path.join(local_dir, "best_model.pt"), device="cuda")  # or "cpu"

See infer.py for the exact image preprocessing (resize to 384×384, normalize to [-1, 1]).

Interactive demo

Run the demo locally:

pip install huggingface_hub
huggingface-cli download kavin-aravindhan/vit-oct-wamd demo_app.py requirements-demo.txt --local-dir .
pip install -r requirements-demo.txt
python demo_app.py

This opens a local web UI where you can upload an OCT B-scan and see the prediction. requirements-demo.txt pins a specific gradio + transformers + huggingface_hub + pydantic/starlette/fastapi combination verified to work together end-to-end — gradio's latest releases require huggingface_hub>=1.0, which conflicts with the transformers==4.53.0 needed to load the checkpoint, and an unpinned newer pydantic/starlette paired with older gradio breaks page rendering outright.

Reproducing training

train.py is a cleaned-up, parameterized version of the exact script used to train this checkpoint — same architecture, loss, hyperparameters (500-trial Optuna search, see training_config.json), and augmentation recipe.

pip install -r requirements-train.txt
python train.py --tfrecord-path /path/to/.tfrecord --output-dir ./runs/my_run

Data access: see Dataset above — the training set (112 labeled images with clinical-finding captions) is hosted as a private HF dataset pending IRB approval.

Note on training-set size: training draws 1000 samples per epoch, with replacement and augmentation, from only 112 unique images — not 1000 unique examples.

Hyperparameter	Value
Image encoder	google/siglip-so400m-patch14-384
Text encoder (auxiliary branch)	google-t5/t5-base
Batch size	8
Learning rate	1e-4
Weight decay	1.6e-6
Alpha (loss mixing, cls vs. alignment)	0.884
Dropout	0.057
Epochs	50 (early stopping patience 20)

License

MIT.

Citation

TODO — paper citation to be added before camera-ready.