initial commit

10 files changed, 682 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..69d6571
--- /dev/null
+++ 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
new file mode 100644
index 0000000..79bf8c0
--- /dev/null
+++ 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
new file mode 100644
index 0000000..696a312
--- /dev/null
+++ 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
new file mode 100644
index 0000000..7a39db4
--- /dev/null
+++ 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
new file mode 100644
index 0000000..6d8b2d9
--- /dev/null
+++ 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
new file mode 100644
index 0000000..38c3cdf
--- /dev/null
+++ 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
new file mode 100644
index 0000000..ab21055
--- /dev/null
+++ 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
new file mode 100644
index 0000000..e77f7b5
--- /dev/null
+++ 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
new file mode 100644
index 0000000..545b299
--- /dev/null
+++ 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
new file mode 100644
index 0000000..a836ca2
--- /dev/null
+++ 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')