# 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.