1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
# EEL4759 Final Project: Diabetic Retinopathy Classification
Compares three pretrained architectures (ResNet-50, EfficientNet-B0, ViT-B/16) with and without CLAHE preprocessing on 5-class DR severity grading.
## Task
Classify fundus photographs into five DR severity grades:
| Label | Grade |
|-------|-------|
| 0 | Healthy |
| 1 | Mild NPDR |
| 2 | Moderate NPDR |
| 3 | Severe NPDR |
| 4 | Proliferative DR |
Dataset: [Diabetic Retinopathy Resized Arranged](https://www.kaggle.com/datasets/amanneo/diabetic-retinopathy-resized-arranged) (auto-downloaded via `kagglehub`).
Split: 70% train / 15% val / 15% test, stratified.
## Experiments
Six experiments (one per model/CLAHE combination):
| Model | CLAHE off | CLAHE on |
|-------|-----------|----------|
| ResNet-50 | `resnet50_clahe=False` | `resnet50_clahe=True` |
| EfficientNet-B0 | `efficientnet_b0_clahe=False` | `efficientnet_b0_clahe=True` |
| ViT-B/16 | `vit_b_16_clahe=False` | `vit_b_16_clahe=True` |
## Design Notes
- CLAHE: applied to the L channel in LAB space before spatial transforms to enhance retinal lesion contrast.
- Class imbalance: `WeightedRandomSampler` with sqrt-inverse-frequency weights plus Focal Loss (gamma=2).
- Discriminative LR: backbone trained at `lr x 0.1`, classification head at `lr`.
- LR schedule: 3-epoch linear warmup followed by cosine annealing.
- Mixed precision: `torch.amp.autocast` + `GradScaler` (CUDA only).
- Early stopping: monitored on validation macro F1, patience=10.
## Files
```
main.py # entry point, experiment loop
dataset.py # CLAHETransform, DRDataset, data loaders
models.py # create_model() for resnet50 / efficientnet_b0 / vit_b_16
train.py # train_model() with AMP, early stopping
evaluate.py # metrics, plots, summary CSV
utils.py # set_seed, get_device, FocalLoss, CLASS_NAMES
requirements.txt # pip dependencies
manifest.scm # Guix environment (AMD ROCm)
```
Each experiment writes to `results/<experiment_key>/`:
- `best_model.pth`
- `classification_report.txt`
- `confusion_matrix.png`
- `training_curves.png`
When more than one experiment finishes, `results/summary.csv` and `results/comparison_chart.png` are written.
## Usage
```bash
# Run all 6 experiments (50 epochs, batch 256, 384px)
python main.py
# Quick sanity check
python main.py --models resnet50 --clahe 0 --epochs 2
# Custom run
python main.py --models resnet50 efficientnet_b0 \
--clahe 0 1 \
--epochs 30 --lr 5e-5 --batch-size 128 \
--img-size 224 --num-workers 4
```
Key arguments:
| Argument | Default | Description |
|----------|---------|-------------|
| `--models` | all three | Models to run |
| `--clahe` | `0 1` | CLAHE variants (0=off, 1=on) |
| `--epochs` | 50 | Max training epochs |
| `--patience` | 10 | Early stopping patience |
| `--batch-size` | 256 | Training batch size |
| `--lr` | 1e-4 | Learning rate (head); backbone gets 10x lower |
| `--img-size` | 384 | Input resolution |
| `--focal-gamma` | 2.0 | Focal loss gamma (0 = standard cross-entropy) |
| `--data-root` | auto | Dataset root; downloaded if omitted |
## Dependencies
```
pip install kagglehub torch torchvision scikit-learn opencv-python-headless matplotlib seaborn
```
For AMD ROCm (Guix): use `manifest.scm`.
## Metrics
Weighted F1, macro F1, accuracy, and quadratic-weighted Cohen's kappa.
|
