Segment blood vessels in retinal fundus images with a PyTorch U-Net (pretrained ResNet34 encoder), a 2-channel CLAHE + Frangi input, trained and evaluated on the DRIVE dataset, and served through a FastAPI endpoint that returns the predicted vessel mask overlaid on the original image — together with a label-free "goodness" score.
Evaluated on a held-out validation split, strictly inside the field-of-view (FOV) mask (the way published DRIVE results are reported):
| Metric | Scratch U-Net (default) | smp + ResNet34 (--arch smp_resnet34) |
|---|---|---|
| F1 (Dice) | 0.833 ✅ | 0.821 ✅ |
| Sensitivity | 0.821 | 0.809 |
| Specificity | 0.977 | 0.975 |
| AUC | 0.980 | 0.976 |
Both use the same 2-channel CLAHE+Frangi input and both clear the commonly cited
> 0.81 DRIVE benchmark. On this small 16-image split the from-scratch U-Net
beat the heavier pretrained model across every patch size tried (64/128px)
and trained more stably — so it's the default. Pick the architecture with
--arch {unet,smp_resnet34}. See Notes & findings.
RGB fundus image
└─ green channel (vessels have highest contrast here)
└─ CLAHE (cv2.createCLAHE, clipLimit=2.0) channel 0
└─ Frangi vesselness (Hessian tubularity filter) channel 1
└─ U-Net (smp ResNet34 encoder, or from-scratch 4+4 blocks)
└─ 1×1 conv → sigmoid → per-pixel vessel probability
- Preprocessing (
dataset.py): green-channel CLAHE stacked with a Frangi vesselness map → 2-channel input; padded to 608×608 (divisible by 32 for the ResNet encoder). - Model (
models.py):build_model()selects either asegmentation_models_pytorchU-Net with an ImageNet-pretrained ResNet34 encoder, or the from-scratch U-Net (unet.py) —Conv → BN → ReLU×2 blocks, MaxPool down, transposed-conv up, skip connections, single-channel output. - Loss (
losses.py):0.5 · BCE + 0.5 · Dice. The Dice term counteracts the severe class imbalance — vessels are only ~8–13% of pixels.
- Data: DRIVE training set (20 annotated images), split 16 train / 4 val.
- Strategy: random 64×64 patch sampling (8000 patches/epoch) whose centres lie inside the FOV. Patch training keeps memory tiny — full-image training on an 8 GB GPU is both slower (~530 s/epoch) and memory-bound. Inputs are cached so the Frangi filter is computed once per image, not every epoch.
- Optimizer: Adam + cosine annealing, flip augmentation, best-validation-F1 checkpoint saved. The pretrained model wants a gentle recipe (lr 1e-4, peaks by ~epoch 15–20); the scratch model uses lr 1e-3 over 150 epochs.
# default: from-scratch U-Net (best & most stable here)
python train.py --arch unet --epochs 120 --lr 1e-3 \
--patch-size 64 --patches-per-epoch 8000 --batch-size 16
# pretrained ResNet34 encoder alternative (gentler recipe)
python train.py --arch smp_resnet34 --epochs 50 --lr 1e-4 \
--patch-size 64 --patches-per-epoch 8000 --batch-size 16
python evaluate.py --checkpoint checkpoints/model_drive.pthuvicorn app:app --port 8000- Open http://localhost:8000/ for a browser upload form, or POST an image:
curl.exe -F "file=@DRIVE/test/images/01_test.tif" \
http://localhost:8000/segment --output overlay.pngPOST /segment→ overlay PNG (vessels in red) with the scores burned in and also returned asX-Vessel-Coverage-Pct/X-Confidence-Pctheaders.POST /segment.json→ just the numeric scores.GET /health→ liveness + device + checkpoint status.
- Vessel coverage % — share of the retina marked as vessel. Healthy fundus images sit around 8–13%; far outside that hints at over/under-segmentation.
- Confidence % — mean decisiveness
max(p, 1−p)of the model inside the FOV, rescaled so 0.5 → 0% and 1.0 → 100%. High means the model rarely sits on the fence. (Typical: ~95%.)
| Original | Segmentation overlay |
|---|---|
 |
|
More overlays in examples/.
DRIVE — Digital Retinal Images for Vessel Extraction. This project expects
the Kaggle distribution placed in DRIVE/:
DRIVE/
├── training/{images, 1st_manual, mask} # 20 images + vessel labels + FOV
└── test/{images, mask} # 20 images + FOV (no public labels)
- Kaggle: https://www.kaggle.com/datasets/andrewmvd/drive-digital-retinal-images-for-vessel-extraction
- Original: https://drive.grand-challenge.org/
Note: the Kaggle DRIVE distribution does not include the second-annotator (
2nd_manual) masks, so inter-annotator agreement is not computed here.
Retinal vessel morphology is a biomarker for diabetic retinopathy, hypertension, and other vascular disease. Automated segmentation supports screening, vessel-density quantification, and as a preprocessing step for downstream diagnosis. This repo is a compact, reproducible reference implementation: classical preprocessing + a clean U-Net + a deployable API.
- Medical image segmentation end-to-end: preprocessing → model → loss → evaluation → deployment.
- Classical + deep fusion: CLAHE and a Hessian-based Frangi vesselness prior stacked as a 2-channel input.
- Two interchangeable models: a from-scratch U-Net and a transfer-learning
U-Net with a pretrained ResNet34 encoder (
segmentation_models_pytorch). - Class-imbalance handling via a combined BCE+Dice objective.
- Correct evaluation: FOV-masked F1/sensitivity/specificity/AUC, matching how the literature reports DRIVE.
- Honest experimentation: a clear, documented finding that the heavier pretrained model did not beat the simpler one on this small dataset.
- Memory-aware training (patch sampling + input caching) that fits an 8 GB consumer GPU.
- Productionization: a FastAPI service with a browser UI, JSON scores, and self-describing overlays.
See IMPROVEMENTS.md for opt-in OpenCV / Hugging Face upgrade
paths.
| File | Role |
|---|---|
dataset.py |
DRIVE loading, CLAHE + Frangi, patch sampling, splits |
models.py |
build_model() — smp ResNet34 or scratch U-Net |
unet.py |
From-scratch U-Net architecture |
losses.py |
BCE+Dice loss, segmentation metrics |
train.py |
Training loop (patch-based, GPU-aware) |
evaluate.py |
FOV-masked F1 / Se / Sp / AUC |
app.py |
FastAPI service + overlay + scores |
IMPROVEMENTS.md |
Optional upgrade paths |
- Pretrained ≠ automatically better. On this 16-image split the heavier smp + ResNet34 model underperformed the from-scratch U-Net (0.833) at every patch size tried: 64px → 0.821 (then overfit-collapses), 128px → 0.817 (stable but lower). Larger patches gave the deep encoder a bigger bottleneck and fixed the instability, but didn't close the accuracy gap. The simpler model wins here — pretrained ImageNet features and a deeper net don't help on 16 domain-specific images, and the bigger net overfits faster.
- Best-checkpoint saving matters: because the pretrained model degrades after its peak, training always keeps the best-validation-F1 snapshot.
- Reproducibility caveat: the val split is only 4 images, so expect ±a few points of variance between runs.
A deliberately honest record of the failures hit building this, and the fixes — because the debugging is half the work:
| Symptom | Root cause | Fix |
|---|---|---|
torch.cuda.is_available() → False on an RTX 5060 |
Installed the CPU-only torch wheel | Reinstall from the CUDA 12.8 index (Blackwell needs cu128, not cu121) |
| Training "produced no output" / looked dead | Python block-buffers stdout when piped; logging only every 5 epochs | -u / flushed per-epoch + per-batch progress lines |
Full-image training: 530 s/epoch then MemoryError |
Full 592²/608² activations thrash host RAM on 8 GB | Patch-based training (48→64px), the standard DRIVE recipe |
| smp run collapsed: val F1 0.82 → 0.33 mid-training | lr 1e-3 over 150 epochs is too aggressive for a pretrained encoder on 16 images | Gentle recipe (lr 1e-4) + best-checkpoint saving; the model peaks by ~epoch 15 |
Recurring Unable to allocate 1.26 MiB (with 4.5 GB free!) |
Not real OOM — a cygwin/bash fork: Resource temporarily unavailable artifact in the shell layer |
Launch training via native PowerShell, not the cygwin shell |
| Frangi recomputed every epoch → slow leak | DriveDataset rebuilt the 2-ch input on every __getitem__ |
Cache inputs once per image |
| Frangi-agreement "confidence" read a misleading 47% | Pearson corr on a very peaky Frangi map is near-zero — an artifact | Dropped the blend; confidence = model decisiveness (~95%) |
| The headline negative: Frangi + pretrained ResNet34 never beat the plain U-Net | Small domain-specific dataset; deeper pretrained net overfits and is starved at the bottleneck | Kept both selectable; shipped the simpler U-Net as default |
pip install -r requirements.txt
# GPU (Blackwell/RTX 50-series needs CUDA 12.8 wheels):
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128MIT — see LICENSE.