commit 8f571ee51854540d186aafc01d25908bba5ea09e
Author: Vineet Kumar <git@vineetk.net>
Date: Sat, 18 Apr 2026 21:34:11 -0400
initial commit
Diffstat:
| A | .gitignore | | | 8 | ++++++++ |
| A | channels.scm | | | 17 | +++++++++++++++++ |
| A | dataset.py | | | 123 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | evaluate.py | | | 146 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | main.py | | | 161 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | manifest.scm | | | 13 | +++++++++++++ |
| A | models.py | | | 35 | +++++++++++++++++++++++++++++++++++ |
| A | requirements.txt | | | 7 | +++++++ |
| A | train.py | | | 136 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | utils.py | | | 37 | +++++++++++++++++++++++++++++++++++++ |
10 files changed, 683 insertions(+), 0 deletions(-)
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,7 @@
+*~
+.claude/
+__pycache__/
+guix-root
+guix-root-1-link
+results/
+venv/
+\ No newline at end of file
diff --git a/channels.scm b/channels.scm
@@ -0,0 +1,17 @@
+(list
+ (channel
+ (name 'guix)
+ (url "https://codeberg.org/dtelsing/Guix")
+ (branch "pytorch"))
+ (channel
+ (name 'guix-rocm)
+ (url "https://codeberg.org/dtelsing/Guix-ROCm")
+ (branch "master")
+ (commit
+ "84964834d67d3e6edeba5bdef7dc42b4413f91c0"))
+ (channel
+ (name 'guix-rocm-ml)
+ (url "https://codeberg.org/dtelsing/Guix-ROCm-ML")
+ (branch "master")
+ (commit
+ "b6c90eccbbb6e280e7222f12102a69aa963d1de4")))
diff --git a/dataset.py b/dataset.py
@@ -0,0 +1,123 @@
+import os
+from PIL import Image
+import numpy as np
+import cv2
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from sklearn.model_selection import train_test_split
+
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+
+class CLAHETransform:
+ """Apply CLAHE to the L channel of LAB color space.
+
+ Operates on PIL Images; returns a PIL Image.
+ """
+
+ def __init__(self, clip_limit: float = 2.0, tile_grid_size: tuple = (8, 8)):
+ self.clip_limit = clip_limit
+ self.tile_grid_size = tile_grid_size
+
+ def __call__(self, img: Image.Image) -> Image.Image:
+ img_np = np.array(img) # RGB uint8
+ lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
+ clahe = cv2.createCLAHE(clipLimit=self.clip_limit, tileGridSize=self.tile_grid_size)
+ lab[:, :, 0] = clahe.apply(lab[:, :, 0])
+ result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
+ return Image.fromarray(result)
+
+
+class DRDataset(Dataset):
+ """Diabetic Retinopathy dataset from flat class-folder layout."""
+
+ def __init__(self, image_paths: list, labels: list, transform=None):
+ self.image_paths = image_paths
+ self.labels = labels
+ self.transform = transform
+
+ def __len__(self) -> int:
+ return len(self.image_paths)
+
+ def __getitem__(self, idx: int):
+ img = Image.open(self.image_paths[idx]).convert('RGB')
+ label = self.labels[idx]
+ if self.transform is not None:
+ img = self.transform(img)
+ return img, label
+
+
+def split_dataset(data_root: str, seed: int = 42):
+ """Scan class folders 0–4 and return stratified 70/15/15 splits.
+
+ Returns:
+ (train_paths, train_labels, val_paths, val_labels, test_paths, test_labels)
+ """
+ all_paths, all_labels = [], []
+ for class_idx in range(5):
+ class_dir = os.path.join(data_root, str(class_idx))
+ for fname in sorted(os.listdir(class_dir)):
+ if fname.lower().endswith(('.jpg', '.jpeg', '.png')):
+ all_paths.append(os.path.join(class_dir, fname))
+ all_labels.append(class_idx)
+
+ # 70% train, 30% temp
+ train_paths, temp_paths, train_labels, temp_labels = train_test_split(
+ all_paths, all_labels, test_size=0.30, stratify=all_labels, random_state=seed
+ )
+ # 50/50 of temp -> 15% val, 15% test
+ val_paths, test_paths, val_labels, test_labels = train_test_split(
+ temp_paths, temp_labels, test_size=0.50, stratify=temp_labels, random_state=seed
+ )
+
+ print(f"Dataset split — train: {len(train_paths)}, val: {len(val_paths)}, test: {len(test_paths)}")
+ return train_paths, train_labels, val_paths, val_labels, test_paths, test_labels
+
+
+def build_transforms(use_clahe: bool, training: bool):
+ """Build a transform pipeline for training or eval."""
+ ops = [transforms.Resize((224, 224))]
+ if use_clahe:
+ ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
+ if training:
+ ops += [
+ transforms.RandomHorizontalFlip(p=0.5),
+ transforms.RandomVerticalFlip(p=0.5),
+ transforms.RandomRotation(degrees=15),
+ ]
+ ops += [
+ transforms.ToTensor(),
+ transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+ ]
+ return transforms.Compose(ops)
+
+
+def build_dataloaders(
+ train_paths, train_labels,
+ val_paths, val_labels,
+ test_paths, test_labels,
+ use_clahe: bool,
+ batch_size: int = 32,
+ num_workers: int = 4,
+):
+ """Create DataLoaders for all three splits."""
+ train_ds = DRDataset(train_paths, train_labels, transform=build_transforms(use_clahe, training=True))
+ val_ds = DRDataset(val_paths, val_labels, transform=build_transforms(use_clahe, training=False))
+ test_ds = DRDataset(test_paths, test_labels, transform=build_transforms(use_clahe, training=False))
+
+ train_loader = DataLoader(
+ train_ds, batch_size=batch_size, shuffle=True,
+ num_workers=num_workers, pin_memory=True
+ )
+ val_loader = DataLoader(
+ val_ds, batch_size=batch_size * 2, shuffle=False,
+ num_workers=num_workers, pin_memory=True
+ )
+ test_loader = DataLoader(
+ test_ds, batch_size=batch_size * 2, shuffle=False,
+ num_workers=num_workers, pin_memory=True
+ )
+ return train_loader, val_loader, test_loader
diff --git a/evaluate.py b/evaluate.py
@@ -0,0 +1,146 @@
+import os
+import csv
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import (
+ accuracy_score,
+ f1_score,
+ classification_report,
+ confusion_matrix,
+ cohen_kappa_score,
+)
+
+from utils import CLASS_NAMES
+
+
+def evaluate_model(model: nn.Module, test_loader: DataLoader, device: torch.device) -> dict:
+ """Run inference on test_loader and return a dict of all metrics."""
+ model.eval()
+ all_preds, all_labels = [], []
+
+ with torch.no_grad():
+ for images, labels in test_loader:
+ images = images.to(device, non_blocking=True)
+ outputs = model(images)
+ preds = outputs.argmax(dim=1)
+ all_preds.extend(preds.cpu().numpy())
+ all_labels.extend(labels.numpy())
+
+ all_preds = np.array(all_preds)
+ all_labels = np.array(all_labels)
+
+ return {
+ 'accuracy': accuracy_score(all_labels, all_preds),
+ 'weighted_f1': f1_score(all_labels, all_preds, average='weighted', zero_division=0),
+ 'macro_f1': f1_score(all_labels, all_preds, average='macro', zero_division=0),
+ 'classification_report': classification_report(
+ all_labels, all_preds, target_names=CLASS_NAMES, zero_division=0
+ ),
+ 'confusion_matrix': confusion_matrix(all_labels, all_preds),
+ 'cohen_kappa': cohen_kappa_score(all_labels, all_preds, weights='quadratic'),
+ }
+
+
+def plot_confusion_matrix(cm: np.ndarray, save_path: str) -> None:
+ """Save a seaborn confusion matrix heatmap to save_path."""
+ fig, ax = plt.subplots(figsize=(8, 6))
+ sns.heatmap(
+ cm, annot=True, fmt='d', cmap='Blues',
+ xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+ ax=ax
+ )
+ ax.set_xlabel('Predicted')
+ ax.set_ylabel('True')
+ ax.set_title('Confusion Matrix')
+ plt.tight_layout()
+ fig.savefig(save_path, dpi=150)
+ plt.close(fig)
+
+
+def plot_training_curves(
+ train_losses: list, val_losses: list, val_f1s: list, save_path: str
+) -> None:
+ """Save loss and F1 training curves to save_path."""
+ epochs = range(1, len(train_losses) + 1)
+ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
+
+ ax1.plot(epochs, train_losses, label='Train loss')
+ ax1.plot(epochs, val_losses, label='Val loss')
+ ax1.set_xlabel('Epoch')
+ ax1.set_ylabel('Loss')
+ ax1.set_title('Training and Validation Loss')
+ ax1.legend()
+
+ ax2.plot(epochs, val_f1s, label='Val weighted F1', color='green')
+ ax2.set_xlabel('Epoch')
+ ax2.set_ylabel('Weighted F1')
+ ax2.set_title('Validation Weighted F1')
+ ax2.legend()
+
+ plt.tight_layout()
+ fig.savefig(save_path, dpi=150)
+ plt.close(fig)
+
+
+def generate_summary(all_results: dict, save_dir: str) -> None:
+ """Print comparison table, save as CSV, and generate bar chart."""
+ os.makedirs(save_dir, exist_ok=True)
+
+ rows = []
+ for key, r in all_results.items():
+ model_name, clahe_str = key.rsplit('_clahe=', 1)
+ rows.append({
+ 'model': model_name,
+ 'clahe': clahe_str,
+ 'weighted_f1': r['weighted_f1'],
+ 'macro_f1': r['macro_f1'],
+ 'accuracy': r['accuracy'],
+ 'cohen_kappa': r['cohen_kappa'],
+ 'train_time_s': round(r['train_time'], 1),
+ 'epochs_trained': r['epochs_trained'],
+ })
+
+ # Print table
+ header = f"{'Model':<18} {'CLAHE':<6} {'W-F1':>6} {'M-F1':>6} {'Acc':>6} {'Kappa':>6} {'Time(s)':>8} {'Epochs':>7}"
+ print('\n' + '=' * len(header))
+ print(header)
+ print('-' * len(header))
+ for row in rows:
+ print(
+ f"{row['model']:<18} {row['clahe']:<6} "
+ f"{row['weighted_f1']:>6.4f} {row['macro_f1']:>6.4f} "
+ f"{row['accuracy']:>6.4f} {row['cohen_kappa']:>6.4f} "
+ f"{row['train_time_s']:>8} {row['epochs_trained']:>7}"
+ )
+ print('=' * len(header))
+
+ # Save CSV
+ csv_path = os.path.join(save_dir, 'summary.csv')
+ with open(csv_path, 'w', newline='') as f:
+ writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+ writer.writeheader()
+ writer.writerows(rows)
+ print(f"\nSummary saved to {csv_path}")
+
+ # Comparison bar chart
+ labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows]
+ w_f1_vals = [r['weighted_f1'] for r in rows]
+
+ fig, ax = plt.subplots(figsize=(10, 5))
+ bars = ax.bar(labels, w_f1_vals, color=['#4C72B0', '#DD8452'] * 3)
+ ax.set_ylabel('Weighted F1-Score')
+ ax.set_title('Weighted F1 Comparison Across Experiments')
+ ax.set_ylim(0, 1.0)
+ for bar, val in zip(bars, w_f1_vals):
+ ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
+ f'{val:.4f}', ha='center', va='bottom', fontsize=9)
+ plt.tight_layout()
+ chart_path = os.path.join(save_dir, 'comparison_chart.png')
+ fig.savefig(chart_path, dpi=150)
+ plt.close(fig)
+ print(f"Comparison chart saved to {chart_path}")
diff --git a/main.py b/main.py
@@ -0,0 +1,161 @@
+"""
+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
+import kagglehub
+
+from utils import set_seed, compute_class_weights, get_device, CLASS_NAMES
+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=30, help='Max training epochs')
+ parser.add_argument('--patience', type=int, default=5, help='Early stopping patience')
+ parser.add_argument('--batch-size', type=int, default=32, 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=4, help='DataLoader workers')
+ 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)
+
+ # Class weights computed once from training labels
+ class_weights = compute_class_weights(train_labels).to(device)
+ print("Class weights:", [f"{w:.3f}" for w in class_weights.cpu().numpy()])
+
+ # ── 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,
+ )
+
+ model = create_model(model_name, num_classes=5, pretrained=True).to(device)
+ criterion = nn.CrossEntropyLoss(weight=class_weights)
+ optimizer = torch.optim.AdamW(
+ model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+ )
+ scheduler = CosineAnnealingLR(optimizer, T_max=args.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
+ ckpt_path = os.path.join(exp_dir, 'best_model.pth')
+ torch.save(train_results['best_model_state'], ckpt_path)
+ print(f" Checkpoint saved: {ckpt_path}")
+
+ # Load best weights for evaluation
+ model.load_state_dict(train_results['best_model_state'])
+
+ # Evaluate on test set
+ test_metrics = evaluate_model(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()
diff --git a/manifest.scm b/manifest.scm
@@ -0,0 +1,13 @@
+(module-define!
+ (resolve-module '(guix-rocm packages rocm-base))
+ '%amdgpu-targets
+ (lambda ()
+ "gfx1100"))
+
+(specifications->manifest
+ (list "python"
+ "python-flask"
+ "python-pytorch-rocm"
+ "python-torchvision-rocm"
+ "rocminfo"
+ "rocm-smi"))
diff --git a/models.py b/models.py
@@ -0,0 +1,35 @@
+import torch.nn as nn
+from torchvision.models import (
+ resnet50, ResNet50_Weights,
+ efficientnet_b0, EfficientNet_B0_Weights,
+ vit_b_16, ViT_B_16_Weights,
+)
+
+
+def create_model(model_name: str, num_classes: int = 5, pretrained: bool = True) -> nn.Module:
+ """Create a pretrained model with the classification head replaced for num_classes.
+
+ Supported model_name values: 'resnet50', 'efficientnet_b0', 'vit_b_16'
+ All backbone layers are trainable (full fine-tuning).
+ """
+ if model_name == 'resnet50':
+ weights = ResNet50_Weights.DEFAULT if pretrained else None
+ model = resnet50(weights=weights)
+ model.fc = nn.Linear(model.fc.in_features, num_classes) # 2048 -> num_classes
+
+ elif model_name == 'efficientnet_b0':
+ weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None
+ model = efficientnet_b0(weights=weights)
+ # classifier is Sequential(Dropout(0.2), Linear(1280, 1000))
+ model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
+
+ elif model_name == 'vit_b_16':
+ weights = ViT_B_16_Weights.DEFAULT if pretrained else None
+ model = vit_b_16(weights=weights)
+ # heads is Sequential(head=Linear(768, 1000))
+ model.heads.head = nn.Linear(model.heads.head.in_features, num_classes)
+
+ else:
+ raise ValueError(f"Unknown model: {model_name!r}. Choose from: resnet50, efficientnet_b0, vit_b_16")
+
+ return model
diff --git a/requirements.txt b/requirements.txt
@@ -0,0 +1,7 @@
+kagglehub
+torch
+torchvision
+scikit-learn
+opencv-python-headless
+matplotlib
+seaborn
diff --git a/train.py b/train.py
@@ -0,0 +1,136 @@
+import copy
+import time
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from sklearn.metrics import f1_score
+
+
+def train_model(
+ model: nn.Module,
+ train_loader: DataLoader,
+ val_loader: DataLoader,
+ criterion: nn.Module,
+ optimizer: torch.optim.Optimizer,
+ scheduler,
+ device: torch.device,
+ num_epochs: int = 30,
+ patience: int = 5,
+ model_name: str = "",
+) -> dict:
+ """Train model with mixed precision, early stopping on val weighted-F1.
+
+ Returns a dict with:
+ best_model_state, train_losses, val_losses, val_f1s,
+ train_time (seconds), epochs_trained
+ """
+ scaler = torch.amp.GradScaler('cuda') if device.type == 'cuda' else None
+
+ best_val_f1 = -1.0
+ best_state = None
+ patience_counter = 0
+
+ train_losses, val_losses, val_f1s = [], [], []
+ start_time = time.time()
+
+ for epoch in range(1, num_epochs + 1):
+ # ── Training phase ──────────────────────────────────────────────────
+ model.train()
+ running_loss = 0.0
+ n_train = 0
+
+ for images, labels in train_loader:
+ images = images.to(device, non_blocking=True)
+ labels = labels.to(device, non_blocking=True)
+
+ optimizer.zero_grad()
+
+ if scaler is not None:
+ with torch.amp.autocast('cuda'):
+ outputs = model(images)
+ loss = criterion(outputs, labels)
+ scaler.scale(loss).backward()
+ scaler.unscale_(optimizer)
+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+ scaler.step(optimizer)
+ scaler.update()
+ else:
+ outputs = model(images)
+ loss = criterion(outputs, labels)
+ loss.backward()
+ torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+ optimizer.step()
+
+ running_loss += loss.item() * images.size(0)
+ n_train += images.size(0)
+
+ train_loss = running_loss / n_train
+ train_losses.append(train_loss)
+
+ # ── Validation phase ─────────────────────────────────────────────────
+ model.eval()
+ val_running_loss = 0.0
+ n_val = 0
+ all_preds, all_labels = [], []
+
+ with torch.no_grad():
+ for images, labels in val_loader:
+ images = images.to(device, non_blocking=True)
+ labels = labels.to(device, non_blocking=True)
+
+ if scaler is not None:
+ with torch.amp.autocast('cuda'):
+ outputs = model(images)
+ loss = criterion(outputs, labels)
+ else:
+ outputs = model(images)
+ loss = criterion(outputs, labels)
+
+ val_running_loss += loss.item() * images.size(0)
+ n_val += images.size(0)
+
+ preds = outputs.argmax(dim=1)
+ all_preds.extend(preds.cpu().numpy())
+ all_labels.extend(labels.cpu().numpy())
+
+ val_loss = val_running_loss / n_val
+ val_losses.append(val_loss)
+
+ val_f1 = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
+ val_f1s.append(val_f1)
+
+ current_lr = scheduler.get_last_lr()[0] if hasattr(scheduler, 'get_last_lr') else optimizer.param_groups[0]['lr']
+ print(
+ f"[{model_name}] Epoch {epoch:3d}/{num_epochs} | "
+ f"train loss: {train_loss:.4f} | val loss: {val_loss:.4f} | "
+ f"val F1: {val_f1:.4f} | lr: {current_lr:.2e}"
+ )
+
+ scheduler.step()
+
+ # ── Early stopping ────────────────────────────────────────────────────
+ if val_f1 > best_val_f1:
+ best_val_f1 = val_f1
+ best_state = copy.deepcopy(model.state_dict())
+ patience_counter = 0
+ else:
+ patience_counter += 1
+ if patience_counter >= patience:
+ print(f" Early stopping at epoch {epoch} (no improvement for {patience} epochs).")
+ break
+
+ train_time = time.time() - start_time
+ epochs_trained = len(train_losses)
+
+ print(f" Training complete: {epochs_trained} epochs, {train_time:.1f}s, best val F1: {best_val_f1:.4f}")
+
+ return {
+ 'best_model_state': best_state,
+ 'train_losses': train_losses,
+ 'val_losses': val_losses,
+ 'val_f1s': val_f1s,
+ 'train_time': train_time,
+ 'epochs_trained': epochs_trained,
+ 'best_val_f1': best_val_f1,
+ }
diff --git a/utils.py b/utils.py
@@ -0,0 +1,37 @@
+import random
+import numpy as np
+import torch
+
+CLASS_NAMES = [
+ 'Healthy (0)',
+ 'Mild NPDR (1)',
+ 'Moderate NPDR (2)',
+ 'Severe NPDR (3)',
+ 'Proliferative DR (4)',
+]
+
+NUM_CLASSES = 5
+
+
+def set_seed(seed: int = 42) -> None:
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+
+
+def compute_class_weights(labels) -> torch.FloatTensor:
+ """Inverse-frequency weighting computed from training labels."""
+ labels = np.array(labels)
+ class_counts = np.bincount(labels, minlength=NUM_CLASSES)
+ total = len(labels)
+ weights = total / (NUM_CLASSES * class_counts.astype(float))
+ return torch.FloatTensor(weights)
+
+
+def get_device() -> torch.device:
+ if torch.cuda.is_available():
+ return torch.device('cuda')
+ return torch.device('cpu')
