Aditya-Sai-19/diabetic-retinopathy-swin

Source: 🤗🏴‍☠️ Rescued by: isiliconimage-classificationapache-2.02M params9 MBsafetensors✓ 2 checksumsupdated 205d ago

Diabetic Retinopathy Prediction — MobileNetV2

This model predicts the severity grade of diabetic retinopathy (DR) from retinal fundus images. It classifies images into 5 grades:

Grade	Label	Description
0	No_DR	No diabetic retinopathy
1	Mild	Mild non-proliferative DR
2	Moderate	Moderate non-proliferative DR
3	Severe	Severe non-proliferative DR
4	Proliferative_DR	Proliferative diabetic retinopathy

Model Details

Architecture: MobileNetV2 (2.2M parameters) fine-tuned for 5-class DR grading
Input: 224×224 RGB retinal fundus images
Training Data: youssefedweqd/Diabetic_Retinopathy_Detection — 25,290 train / 2,810 val / 7,026 test images
Loss Function: Focal Loss (γ=2.0) with inverse-frequency class weights — critical for handling severe class imbalance
Primary Metric: Quadratic Weighted Kappa (QWK) — the standard metric for ordinal DR grading

Training Recipe

Based on published DR detection literature (see References below):

Parameter	Value
Base model	google/mobilenet_v2_1.0_224 (ImageNet pretrained)
Optimizer	AdamW (lr=3e-4, weight_decay=0.01)
LR Schedule	Cosine with 200 warmup steps
Batch size	32
Epochs	5
Loss	Focal Loss (γ=2.0) + class weights
Augmentations	RandomResizedCrop, H/V flips, rotation ±30°, color jitter
Hardware	CPU (2 vCPU, 16GB RAM)

Class Distribution (severe imbalance)

Grade	Count	%
0 - No DR	18,583	73.5%
1 - Mild	1,759	7.0%
2 - Moderate	3,811	15.1%
3 - Severe	628	2.5%
4 - Proliferative	509	2.0%

Results

Validation Performance (best epoch = 3)

Epoch	Loss	Accuracy	QWK
1	1.064	9.0%	0.182
2	0.819	11.6%	0.345
3	0.828	13.6%	0.373
4	0.889	14.0%	0.298
5	0.775	14.2%	0.363

Test Set Performance

Accuracy: 13.5%
QWK: 0.384

Per-Class Test Results

Class	Precision	Recall	F1
No DR	0.00	0.00	0.00
Mild	0.08	0.90	0.14
Moderate	0.43	0.31	0.36
Severe	0.34	0.55	0.42
Proliferative	0.32	0.63	0.42

Usage

from transformers import pipeline

classifier = pipeline("image-classification", model="Aditya-Sai-19/diabetic-retinopathy-swin")
result = classifier("path/to/retinal_image.jpg")
print(result)
# [{'label': 'Moderate', 'score': 0.45}, {'label': 'Mild', 'score': 0.22}, ...]

Or manually:

from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

processor = AutoImageProcessor.from_pretrained("Aditya-Sai-19/diabetic-retinopathy-swin")
model = AutoModelForImageClassification.from_pretrained("Aditya-Sai-19/diabetic-retinopathy-swin")

image = Image.open("retinal_scan.jpg")
inputs = processor(image, return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits
    predicted_class = logits.argmax(-1).item()

print(f"Predicted: {model.config.id2label[predicted_class]}")

Limitations & Recommendations for Improvement

This model was trained on CPU only (free tier) with a lightweight MobileNetV2 backbone. For production-quality DR grading, consider:

Use a larger backbone: Swin Transformer Base or EfficientNet-B4 (published SOTA: QWK > 0.93 on APTOS)
Train on GPU: Enables larger batch sizes, more epochs, and larger models
Ben Graham preprocessing: Apply contrast enhancement (I_ce = 4*I - 4*G(σ=w/90)*I + 128) and FOV cropping — this alone adds +2-5% AUC
More data: Merge APTOS + Messidor-2 + IDRiD + EyePACS for better generalization
Data augmentation: 10x oversampling of minority classes (grades 3 & 4)
Post-processing: Optimized ordinal thresholds for QWK maximization

References

Dual Branch Deep Learning Network for Detection and Stage Grading of DR (arXiv:2308.09945) — QWK=0.930
A ResNet is All You Need? (arXiv:2210.03180) — AUC=0.955
Ord2Seq: Ordinal Regression as Label Sequence Prediction (arXiv:2307.09004) — Acc=84.2%

Framework Versions

Transformers 5.6.2
PyTorch 2.11.0
Datasets 4.8.4