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