EyePACS Diabetic Retinopathy Grading

Ordinal multiple-instance learning (MIL) for 5-class diabetic-retinopathy (DR) grading from fundus photographs (EyePACS / Kaggle DR scale, grades 0–4: No DR, Mild, Moderate, Severe, Proliferative).

The model treats DR grading as ordered lesion-evidence accumulation over multi-scale crops of a fundus image, rather than as flat 5-way classification. It ships a baseline attention-MIL model, a threshold-specific ordinal MIL model, and a selective-aggregator ordinal MIL baseline, all behind one config-driven pipeline.

---

Technical approach

1. Preprocessing (src/bloom_dr/data/preprocessing.py)

Each image is turned into a small bag of views:

1. Black-border crop — trims the dark frame around the circular fundus.
2. Ben Graham enhancement (optional) — 4·img − 4·blur + 128, which normalises illumination and sharpens vessels / micro-lesions.
3. Multi-view cropping — one global resized view plus several local crops. Two crop policies:
   • random — random rectangular crops at configured scales.
   • lesion_candidate — heuristic mining of red-lesion (dark-in-green-channel) and bright-lesion (locally bright) regions via a local-contrast saliency map, with spatial-diversity suppression. This is a heuristic prior, not a trained lesion detector. Boxes are kept for auditing; see scripts/visualize_lesion_crops.py.
4. Normalisation to ImageNet statistics; output tensor is [V, 3, H, W].

2. Encoder + aggregation (src/bloom_dr/models/)

views [B,V,3,H,W] → shared backbone → per-crop embeddings [B,V,D] → MIL pooling → heads

• Backbone — small_cnn (dependency-free, for CPU/dev) or any timm:<name> model (e.g. timm:tf_efficientnet_b0_ns).

• Aggregation — three architectures (selectable via model.architecture):

  architecture	Aggregation
  baseline	single gated-attention MIL query → one pooled vector → ordinal head
  tom	threshold-specific gated-attention MIL: a separate attention query per ordinal boundary P(y≥1..4), so evidence aggregation is boundary-aware
  selective_aggregator	transformer-style aggregator token per ordinal boundary (an in-repo reimplementation of selective-aggregator ordinal MIL, Shiku et al., WACV 2025) used as a comparison baseline

• Heads — ordinal head (P(y≥k) logits), auxiliary 5-way classification head, per-crop ordinal head (auxiliary), and an optional dedicated grade-0-vs-grade-1 "mild" head.

3. Ordinal modelling and monotonicity (src/bloom_dr/evaluation/metrics.py)

The model predicts cumulative probabilities P(y≥1) … P(y≥4). At inference these are projected to be rank-monotone (P(y≥1) ≥ … ≥ P(y≥4)) via a running minimum, then converted to a proper class distribution. A soft monotonicity penalty is also applied during training, and the monotonic-violation rate is reported each epoch.

4. Loss (src/bloom_dr/training/losses.py)

A weighted combination (each term toggle-able by weight):

• Ordinal BCE over all thresholds (class-imbalance pos_weight).
• Cross-entropy on the auxiliary class head (inverse-frequency class weights).
• Monotonic penalty — Σ ReLU(P(y≥k+1) − P(y≥k))².
• Mild rescue — BCE on the dedicated grade-0/1 head (or threshold-0 logit), restricted to the normal/mild subset, targeting the grade-0/1 boundary that is the hardest and most clinically important.
• Mild margin — ranking term pushing mild scores above normal scores.
• Attention diversity — discourages threshold-specific attention queries from collapsing to the same distribution.
• Crop-ordinal (optional, off by default).

5. Imbalance handling

DR is heavily imbalanced (~73% grade 0). Two levers:

• inverse-frequency class / pos_weight weighting in the loss, and
• an optional mild_oversample sampler that rebalances classes and additionally boosts grade 1, with the resulting class mix logged transparently.

6. Evaluation protocol (src/bloom_dr/training/trainer.py)

Patients (not images) are split whole and stratified by their most-severe grade, preventing left/right-eye leakage. Four patient-disjoint partitions:

Partition	Role
train	gradient updates
val	epoch / checkpoint selection
calib	threshold (and future temperature) fitting only
test	final metrics, computed once

Operating thresholds on the ordinal severity score are fit on calib (via a deterministic grid coordinate-ascent that maximises QWK), the best epoch is selected on val, and final numbers are reported once on test. The trainer asserts the partitions are disjoint and flags a run as degraded in summary.json if calib /test are not configured.

7. Metrics

Quadratic Weighted Kappa (QWK), macro-F1, per-grade F1 (incl. a named grade-1 F1), per-grade recall, mean absolute error, severe-error rate (|err|≥2), referable-DR AUC from P(y≥2), severe-DR AUC from P(y≥3), ECE and Brier score (computed on the deployed ordinal decision), monotonic-violation rate, and patient-cluster bootstrap confidence intervals.

---

Setup

pip install -e .                 # core: numpy, pandas, pillow, pyyaml, torch, matplotlib
pip install -e ".[research]"     # optional: timm (real backbones), scipy, kagglehub

small_cnn runs without timm. Real runs use a timm: backbone (needs the research extra).

Data layout

<DATA_ROOT>/
  trainLabels.csv          # columns: image,level   (level in 0..4)
  train/<image>.jpeg
  test/<image>.jpeg        # optional (Kaggle submission)
  sampleSubmission.csv     # optional

Patient IDs are the prefix before _, so 10_left and 10_right always stay in the same partition. Configs read ${BLOOM_DATA_ROOT:-/workspace/eyepacs-dataset}; set BLOOM_DATA_ROOT / BLOOM_OUTPUT_DIR to relocate inputs / outputs.

---

Usage

Train:

python -m bloom_dr.train --config configs/ablations/A5_tom_full.yaml

Available configs include configs/sample*.yaml (tiny CPU smoke), configs/runpod*.yaml (full GPU runs), and the controlled ablation grid in configs/ablations/A0…A5.yaml (baseline → selective-aggregator → TOM → +lesion crops → +mild rescue → full).

Evaluate a checkpoint on the held-out split:

python -m bloom_dr.evaluate --config <cfg> --checkpoint <run>/best.pt --output-dir <run>/eval

Evaluate on an external dataset (APTOS / MESSIDOR-2 / custom; refuses to run if the data is absent, and reuses the internal calibration thresholds):

python -m bloom_dr.evaluate_external --config <cfg> --checkpoint <run>/best.pt \
  --dataset aptos --labels-csv /data/aptos/train.csv --image-dir /data/aptos/train_images \
  --output-dir <run>/external_aptos

Predict a Kaggle submission, generate figures, visualise lesion crops:

python -m bloom_dr.predict  --config <cfg> --checkpoint <run>/best.pt --output <run>/submission.csv
python -m bloom_dr.figures  --labels-csv <DATA_ROOT>/trainLabels.csv --run-dir <run> --output-dir <run>/figures
python scripts/visualize_lesion_crops.py --image-dir <DATA_ROOT>/train --output-dir <run>/lesion_debug

Tests:

pytest -q

---

Repository layout

src/bloom_dr/
  config.py                 # dataclass configs + YAML loader
  data/                     # splits, dataset, preprocessing, external adapters
  models/                   # backbones, MIL pooling, heads, model assembly
  training/                 # trainer (protocol), losses
  evaluation/               # metrics (QWK, AUC, ECE/Brier, monotonicity, bootstrap)
  visualization/            # paper figures
  observability/, monarch/  # auxiliary experiment modules
  train.py / evaluate.py / evaluate_external.py / predict.py / figures.py / explain.py
configs/                    # sample, runpod, and ablation YAMLs
scripts/                    # smoke test, lesion-crop viz, RunPod tooling
tests/                      # split / metrics / model / loss / external / e2e smoke

Notes

• The lesion_candidate crop policy is a heuristic saliency prior, not clinical lesion detection.
• Ordinal monotonicity is enforced at inference (projection) plus a soft training penalty, not architecturally guaranteed.
• The selective_aggregator model is an in-repo reimplementation of WACV-2025 selective-aggregator ordinal MIL, provided for comparison, not the official code.