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.

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

If you use this model or the associated dataset in your research, please cite the following paper:

Co-Annotator: Expert-Distilled ViT and VLM for Visual and Documentation Guidance in Age-Related Macular Degeneration

ACM Symposium on User Interface Software and Technology (UIST '26)

@inproceedings{li2026coannotator,
  title     = {Co-Annotator: Expert-Distilled ViT and VLM for Visual and Documentation Guidance in Age-Related Macular Degeneration},
  author    = {Li, Ziheng and Freeman, Benjamin and Raman, Akshay and Rajkumar, Kavin Aravindhan and Fang, Xinxin and Srivastava, Rishabh and Feiner, Steven and Thakoor, Kaveri A.},
  booktitle = {Proceedings of the 39th Annual ACM Symposium on User Interface Software and Technology (UIST '26)},
  year      = {2026},
  publisher = {Association for Computing Machinery},
  doi       = {10.1145/3830398.3830722},
  url       = {https://doi.org/10.1145/3830398.3830722}
}

Paper: arXiv | ACM DOI