Prediction of Incident Atrial Fibrillation from Retinal Fundus Images Using a Multimodal Foundation Model

This repository contains the minimum publication code for training, inference, evaluation, and figure drawing for the proposed biomarkers:

• RetiAF: retinal fundus photograph model
• Hybrid_RetiAF: retinal fundus photograph plus CHADSVAS clinical score model

The code is adapted from the RETFound fine-tuning workflow used in the study. Data, protected health information, experiment logs, and trained checkpoints are not included.

Repository Structure

.
  train.py                    # Fine-tune RetiAF or Hybrid_RetiAF
  infer.py                    # Generate per-image prediction scores
  figures/draw_figures.py     # Draw ROC, PR, and calibration figures
  src/
    data.py                   # Image and clinical-feature loading
    metrics.py                # Binary classification metrics
    modeling.py               # Model creation and checkpoint loading
    models_vit.py             # RETFound ViT-L/16 and fused clinical model
    plotting.py               # Shared plotting utilities
  examples/
    clinical_schema.csv
    predictions_schema.csv
  checkpoints/
    README.md

Installation

The original experiments used PyTorch 1.8.1, CUDA 11.1, and timm==0.3.2. A newer compatible PyTorch build can also work, but timm==0.3.2 should be retained unless model-loading behavior is revalidated.

conda create -n retiaf python=3.9 -y
conda activate retiaf
pip install -r requirements.txt

Install the CUDA-specific PyTorch wheel appropriate for your system if the generic command does not match your GPU environment.

Data Layout

Training and labelled evaluation data should use ImageFolder-style class folders:

data/
  train/
    class_0/
      image001.png
    class_1/
      image002.png
  val/
    class_0/
    class_1/
  test/
    class_0/
    class_1/

class_0 means no incident AF and class_1 means incident AF.

For Hybrid_RetiAF, provide a CSV with image identifiers and CHADSVAS:

image_name,CHADSVAS
image001.png,1
image002.png,3

Accepted identifier columns include image_name, Patient, patient_id, filename, file_name, and eid.

If a CHADSVAS column is absent, the code can calculate the score from available raw variables: age, sex/female indicator, heart failure, hypertension, prior stroke/TIA/thromboembolism, diabetes, and peripheral arterial disease. Missing binary values are treated as 0; missing age is mean-imputed within the supplied CSV.

Training

Fine-tune the image-only model from RETFound CFP weights:

python train.py \
  --data-dir /path/to/data \
  --output-dir runs/retiaf \
  --model retiaf \
  --pretrained checkpoints/RETFound_cfp_weights.pth \
  --epochs 100 \
  --batch-size 16 \
  --lr 1e-3 \
  --weight-decay 0.01 \
  --input-size 224 \
  --device cuda

Fine-tune the hybrid model:

python train.py \
  --data-dir /path/to/data \
  --clinical-csv /path/to/clinical.csv \
  --output-dir runs/hybrid_retiaf \
  --model hybrid_retiaf \
  --pretrained checkpoints/RETFound_cfp_weights.pth \
  --epochs 100 \
  --batch-size 16 \
  --lr 1e-3 \
  --weight-decay 0.01 \
  --input-size 224 \
  --device cuda

Outputs include:

• checkpoint-best.pth
• training_log.csv
• val_predictions_latest.csv
• test_predictions.csv
• metrics_test.json

Inference

For labelled data, folders can remain in class_0 and class_1. For unlabelled data, a single folder of images is also accepted:

unlabelled_images/
  image001.png
  image002.png

Generate both biomarkers:

python infer.py \
  --input-dir /path/to/images_or_imagefolder_dataset \
  --clinical-csv /path/to/clinical.csv \
  --retiaf-checkpoint checkpoints/retiaf_checkpoint-best.pth \
  --hybrid-checkpoint checkpoints/hybrid_retiaf_checkpoint-best.pth \
  --output-dir results/inference \
  --batch-size 16 \
  --device cuda

For unlabelled data, individual_predictions.csv contains per-image scores for downstream linkage. AUROC, accuracy, sensitivity, specificity, and PPV are only calculated when labels are available for all images and both classes are present.

Figure Drawing

Draw ROC, precision-recall, calibration figures, and a metrics table from a labelled prediction CSV:

python figures/draw_figures.py \
  --predictions results/inference/individual_predictions.csv \
  --output-dir results/figures \
  --label-column gt_I48 \
  --score-columns RetiAF_Score Hybrid_RetiAF_Score

Outputs include:

• roc_curves.png and .svg
• precision_recall_curves.png and .svg
• calibration.png and .svg
• figure_metrics.csv
• figure_metrics.json

Image Preprocessing

Evaluation preprocessing follows the project setting:

Bicubic resize to 256 -> center crop to 224 -> ImageNet normalization

The training transform uses ImageNet-style augmentation from timm.

Checkpoints

Place pretrained and trained checkpoints under checkpoints/ or pass absolute paths. Large checkpoint files are ignored by git by default.