""" EEL4759 Final Project — Diabetic Retinopathy Classification Compares ResNet-50, EfficientNet-B0, and ViT-B/16 with and without CLAHE preprocessing. Usage: python main.py # run all 6 experiments (default) python main.py --models resnet50 --clahe 0 --epochs 2 # quick sanity check """ import argparse import os import torch import torch.nn as nn from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR import kagglehub from utils import set_seed, get_device, FocalLoss from dataset import split_dataset, build_dataloaders from models import create_model from train import train_model from evaluate import evaluate_model, plot_confusion_matrix, plot_training_curves, generate_summary RESULTS_DIR = 'results' def parse_args(): parser = argparse.ArgumentParser(description='Diabetic Retinopathy Classification') parser.add_argument( '--models', nargs='+', default=['resnet50', 'efficientnet_b0', 'vit_b_16'], help='Models to run' ) parser.add_argument( '--clahe', nargs='+', type=int, default=[0, 1], help='CLAHE settings to run (0=off, 1=on)' ) parser.add_argument('--epochs', type=int, default=50, help='Max training epochs') parser.add_argument('--patience', type=int, default=10, help='Early stopping patience') parser.add_argument('--batch-size', type=int, default=256, help='Training batch size') parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') parser.add_argument('--weight-decay', type=float, default=1e-4, help='AdamW weight decay') parser.add_argument('--seed', type=int, default=42, help='Random seed') parser.add_argument('--data-root', type=str, default=None, help='Path to dataset root (auto-downloaded if not set)') parser.add_argument('--num-workers', type=int, default=8, help='DataLoader workers') parser.add_argument('--focal-gamma', type=float, default=2.0, help='Focal loss gamma (0=standard CE)') parser.add_argument('--warmup-epochs', type=int, default=3, help='LR warmup epochs') parser.add_argument('--img-size', type=int, default=384, help='Input image size') return parser.parse_args() def main(): args = parse_args() set_seed(args.seed) device = get_device() print(f"Using device: {device}") # ── Dataset ────────────────────────────────────────────────────────────── if args.data_root is None: print("Downloading/locating dataset via kagglehub...") data_root = kagglehub.dataset_download("amanneo/diabetic-retinopathy-resized-arranged") else: data_root = args.data_root print(f"Dataset root: {data_root}") train_paths, train_labels, val_paths, val_labels, test_paths, test_labels = \ split_dataset(data_root, seed=args.seed) # ── Experiments ────────────────────────────────────────────────────────── experiments = [ (model_name, bool(use_clahe)) for model_name in args.models for use_clahe in args.clahe ] all_results = {} os.makedirs(RESULTS_DIR, exist_ok=True) for model_name, use_clahe in experiments: exp_key = f"{model_name}_clahe={use_clahe}" exp_dir = os.path.join(RESULTS_DIR, exp_key) os.makedirs(exp_dir, exist_ok=True) print(f"\n{'='*70}") print(f"Experiment: {model_name} | CLAHE: {use_clahe}") print(f"{'='*70}") train_loader, val_loader, test_loader = build_dataloaders( train_paths, train_labels, val_paths, val_labels, test_paths, test_labels, use_clahe=use_clahe, batch_size=args.batch_size, num_workers=args.num_workers, img_size=args.img_size, ) model = create_model(model_name, num_classes=5, pretrained=True).to(device) # Discriminative LR: backbone gets 10x lower LR than the new head if model_name == 'resnet50': head_params = list(model.fc.parameters()) elif model_name == 'efficientnet_b0': head_params = list(model.classifier.parameters()) elif model_name == 'vit_b_16': head_params = list(model.heads.parameters()) else: head_params = [] head_ids = {id(p) for p in head_params} backbone_params = [p for p in model.parameters() if id(p) not in head_ids] if torch.cuda.device_count() > 1: print(f" Using {torch.cuda.device_count()} GPUs via DataParallel") model = torch.nn.DataParallel(model) criterion = FocalLoss(alpha=None, gamma=args.focal_gamma) optimizer = torch.optim.AdamW( [ {'params': backbone_params, 'lr': args.lr * 0.1}, {'params': head_params, 'lr': args.lr}, ], weight_decay=args.weight_decay, ) warmup_epochs = min(args.warmup_epochs, args.epochs - 1) warmup_scheduler = LinearLR(optimizer, start_factor=0.1, total_iters=warmup_epochs) cosine_scheduler = CosineAnnealingLR(optimizer, T_max=max(args.epochs - warmup_epochs, 1)) scheduler = SequentialLR(optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_epochs]) # Train train_results = train_model( model, train_loader, val_loader, criterion, optimizer, scheduler, device, num_epochs=args.epochs, patience=args.patience, model_name=exp_key, ) # Save best checkpoint (strip DataParallel 'module.' prefix for portability) ckpt_path = os.path.join(exp_dir, 'best_model.pth') best_state = train_results['best_model_state'] if isinstance(model, torch.nn.DataParallel): best_state = {k[7:]: v for k, v in best_state.items()} torch.save(best_state, ckpt_path) print(f" Checkpoint saved: {ckpt_path}") # Load best weights for evaluation core_model = model.module if isinstance(model, torch.nn.DataParallel) else model core_model.load_state_dict(best_state) # Evaluate on test set test_metrics = evaluate_model(core_model, test_loader, device) print(f"\n Test results — Weighted F1: {test_metrics['weighted_f1']:.4f} | " f"Macro F1: {test_metrics['macro_f1']:.4f} | " f"Accuracy: {test_metrics['accuracy']:.4f} | " f"Kappa: {test_metrics['cohen_kappa']:.4f}") print(test_metrics['classification_report']) # Save per-class report report_path = os.path.join(exp_dir, 'classification_report.txt') with open(report_path, 'w') as f: f.write(f"Experiment: {exp_key}\n\n") f.write(test_metrics['classification_report']) # Plots plot_confusion_matrix( test_metrics['confusion_matrix'], save_path=os.path.join(exp_dir, 'confusion_matrix.png'), ) plot_training_curves( train_results['train_losses'], train_results['val_losses'], train_results['val_f1s'], save_path=os.path.join(exp_dir, 'training_curves.png'), ) all_results[exp_key] = {**train_results, **test_metrics} # Remove bulky state dict from summary dict (already saved to disk) all_results[exp_key].pop('best_model_state', None) # ── Summary ─────────────────────────────────────────────────────────────── if len(all_results) > 1: generate_summary(all_results, save_dir=RESULTS_DIR) if __name__ == '__main__': main()