main.py (8032B)
1 """ 2 EEL4759 Final Project — Diabetic Retinopathy Classification 3 Compares ResNet-50, EfficientNet-B0, and ViT-B/16 with and without CLAHE preprocessing. 4 5 Usage: 6 python main.py # run all 6 experiments (default) 7 python main.py --models resnet50 --clahe 0 --epochs 2 # quick sanity check 8 """ 9 10 import argparse 11 import os 12 13 import torch 14 import torch.nn as nn 15 from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR 16 import kagglehub 17 18 from utils import set_seed, get_device, FocalLoss 19 from dataset import split_dataset, build_dataloaders 20 from models import create_model 21 from train import train_model 22 from evaluate import evaluate_model, plot_confusion_matrix, plot_training_curves, generate_summary 23 24 RESULTS_DIR = 'results' 25 26 27 def parse_args(): 28 parser = argparse.ArgumentParser(description='Diabetic Retinopathy Classification') 29 parser.add_argument( 30 '--models', nargs='+', 31 default=['resnet50', 'efficientnet_b0', 'vit_b_16'], 32 help='Models to run' 33 ) 34 parser.add_argument( 35 '--clahe', nargs='+', type=int, default=[0, 1], 36 help='CLAHE settings to run (0=off, 1=on)' 37 ) 38 parser.add_argument('--epochs', type=int, default=50, help='Max training epochs') 39 parser.add_argument('--patience', type=int, default=10, help='Early stopping patience') 40 parser.add_argument('--batch-size', type=int, default=256, help='Training batch size') 41 parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') 42 parser.add_argument('--weight-decay', type=float, default=1e-4, help='AdamW weight decay') 43 parser.add_argument('--seed', type=int, default=42, help='Random seed') 44 parser.add_argument('--data-root', type=str, default=None, 45 help='Path to dataset root (auto-downloaded if not set)') 46 parser.add_argument('--num-workers', type=int, default=8, help='DataLoader workers') 47 parser.add_argument('--focal-gamma', type=float, default=2.0, help='Focal loss gamma (0=standard CE)') 48 parser.add_argument('--warmup-epochs', type=int, default=3, help='LR warmup epochs') 49 parser.add_argument('--img-size', type=int, default=384, help='Input image size') 50 return parser.parse_args() 51 52 53 def main(): 54 args = parse_args() 55 set_seed(args.seed) 56 device = get_device() 57 print(f"Using device: {device}") 58 59 # ── Dataset ────────────────────────────────────────────────────────────── 60 if args.data_root is None: 61 print("Downloading/locating dataset via kagglehub...") 62 data_root = kagglehub.dataset_download("amanneo/diabetic-retinopathy-resized-arranged") 63 else: 64 data_root = args.data_root 65 print(f"Dataset root: {data_root}") 66 67 train_paths, train_labels, val_paths, val_labels, test_paths, test_labels = \ 68 split_dataset(data_root, seed=args.seed) 69 70 # ── Experiments ────────────────────────────────────────────────────────── 71 experiments = [ 72 (model_name, bool(use_clahe)) 73 for model_name in args.models 74 for use_clahe in args.clahe 75 ] 76 77 all_results = {} 78 os.makedirs(RESULTS_DIR, exist_ok=True) 79 80 for model_name, use_clahe in experiments: 81 exp_key = f"{model_name}_clahe={use_clahe}" 82 exp_dir = os.path.join(RESULTS_DIR, exp_key) 83 os.makedirs(exp_dir, exist_ok=True) 84 85 print(f"\n{'='*70}") 86 print(f"Experiment: {model_name} | CLAHE: {use_clahe}") 87 print(f"{'='*70}") 88 89 train_loader, val_loader, test_loader = build_dataloaders( 90 train_paths, train_labels, 91 val_paths, val_labels, 92 test_paths, test_labels, 93 use_clahe=use_clahe, 94 batch_size=args.batch_size, 95 num_workers=args.num_workers, 96 img_size=args.img_size, 97 ) 98 99 model = create_model(model_name, num_classes=5, pretrained=True).to(device) 100 101 # Discriminative LR: backbone gets 10x lower LR than the new head 102 if model_name == 'resnet50': 103 head_params = list(model.fc.parameters()) 104 elif model_name == 'efficientnet_b0': 105 head_params = list(model.classifier.parameters()) 106 elif model_name == 'vit_b_16': 107 head_params = list(model.heads.parameters()) 108 else: 109 head_params = [] 110 head_ids = {id(p) for p in head_params} 111 backbone_params = [p for p in model.parameters() if id(p) not in head_ids] 112 113 if torch.cuda.device_count() > 1: 114 print(f" Using {torch.cuda.device_count()} GPUs via DataParallel") 115 model = torch.nn.DataParallel(model) 116 117 criterion = FocalLoss(alpha=None, gamma=args.focal_gamma) 118 optimizer = torch.optim.AdamW( 119 [ 120 {'params': backbone_params, 'lr': args.lr * 0.1}, 121 {'params': head_params, 'lr': args.lr}, 122 ], 123 weight_decay=args.weight_decay, 124 ) 125 warmup_epochs = min(args.warmup_epochs, args.epochs - 1) 126 warmup_scheduler = LinearLR(optimizer, start_factor=0.1, total_iters=warmup_epochs) 127 cosine_scheduler = CosineAnnealingLR(optimizer, T_max=max(args.epochs - warmup_epochs, 1)) 128 scheduler = SequentialLR(optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_epochs]) 129 130 # Train 131 train_results = train_model( 132 model, train_loader, val_loader, 133 criterion, optimizer, scheduler, 134 device, 135 num_epochs=args.epochs, 136 patience=args.patience, 137 model_name=exp_key, 138 ) 139 140 # Save best checkpoint (strip DataParallel 'module.' prefix for portability) 141 ckpt_path = os.path.join(exp_dir, 'best_model.pth') 142 best_state = train_results['best_model_state'] 143 if isinstance(model, torch.nn.DataParallel): 144 best_state = {k[7:]: v for k, v in best_state.items()} 145 torch.save(best_state, ckpt_path) 146 print(f" Checkpoint saved: {ckpt_path}") 147 148 # Load best weights for evaluation 149 core_model = model.module if isinstance(model, torch.nn.DataParallel) else model 150 core_model.load_state_dict(best_state) 151 152 # Evaluate on test set 153 test_metrics = evaluate_model(core_model, test_loader, device) 154 155 print(f"\n Test results — Weighted F1: {test_metrics['weighted_f1']:.4f} | " 156 f"Macro F1: {test_metrics['macro_f1']:.4f} | " 157 f"Accuracy: {test_metrics['accuracy']:.4f} | " 158 f"Kappa: {test_metrics['cohen_kappa']:.4f}") 159 print(test_metrics['classification_report']) 160 161 # Save per-class report 162 report_path = os.path.join(exp_dir, 'classification_report.txt') 163 with open(report_path, 'w') as f: 164 f.write(f"Experiment: {exp_key}\n\n") 165 f.write(test_metrics['classification_report']) 166 167 # Plots 168 plot_confusion_matrix( 169 test_metrics['confusion_matrix'], 170 save_path=os.path.join(exp_dir, 'confusion_matrix.png'), 171 ) 172 plot_training_curves( 173 train_results['train_losses'], 174 train_results['val_losses'], 175 train_results['val_f1s'], 176 save_path=os.path.join(exp_dir, 'training_curves.png'), 177 ) 178 179 all_results[exp_key] = {**train_results, **test_metrics} 180 # Remove bulky state dict from summary dict (already saved to disk) 181 all_results[exp_key].pop('best_model_state', None) 182 183 # ── Summary ─────────────────────────────────────────────────────────────── 184 if len(all_results) > 1: 185 generate_summary(all_results, save_dir=RESULTS_DIR) 186 187 188 if __name__ == '__main__': 189 main()