---
license: mit
base_model: google/siglip-so400m-patch14-384
tags:
  - medical-imaging
  - ophthalmology
  - oct
  - age-related-macular-degeneration
  - vision-transformer
  - siglip
  - pytorch
pipeline_tag: image-classification
---

# vit-oct-wamd

A [SigLIP](https://huggingface.co/google/siglip-so400m-patch14-384) 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](https://huggingface.co/datasets/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

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

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

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

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

**Data access:** see [Dataset](#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.