eel4759_classification

Comparison of different ImageNet-based CNN models for classifying diabetic retinopathy images
Log | Files | Refs | README

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.