5 files changed, 160 insertions, 63 deletions

diff --git a/dataset.py b/dataset.py
index 696a312..0068a97 100644
--- a/dataset.py
+++ b/dataset.py
@@ -4,7 +4,7 @@ import numpy as np
 import cv2
 
 import torch
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
 from torchvision import transforms
 from sklearn.model_selection import train_test_split
 
@@ -19,14 +19,12 @@ class CLAHETransform:
     """
 
     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
+        self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=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])
+        lab[:, :, 0] = self.clahe.apply(lab[:, :, 0])
         result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
         return Image.fromarray(result)
 
@@ -77,17 +75,32 @@ def split_dataset(data_root: str, seed: int = 42):
     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)))
+def build_transforms(use_clahe: bool, training: bool, img_size: int = 224):
+    """Build a transform pipeline for training or eval.
+
+    Training path uses RandomResizedCrop for scale/position augmentation.
+    Eval path uses a deterministic Resize to preserve comparability.
+    CLAHE is applied before spatial transforms at a slightly larger size
+    so the crop has room to operate.
+    """
+    ops = []
     if training:
+        if use_clahe:
+            # Resize larger so RandomResizedCrop still sees enough context after CLAHE
+            ops.append(transforms.Resize((img_size + 32, img_size + 32)))
+            ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
         ops += [
+            transforms.RandomResizedCrop(img_size, scale=(0.8, 1.0), ratio=(0.9, 1.1)),
             transforms.RandomHorizontalFlip(p=0.5),
             transforms.RandomVerticalFlip(p=0.5),
-            transforms.RandomRotation(degrees=15),
+            transforms.RandomRotation(degrees=30),
+            transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.02),
+            transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 1.0)),
         ]
+    else:
+        ops.append(transforms.Resize((img_size, img_size)))
+        if use_clahe:
+            ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
     ops += [
         transforms.ToTensor(),
         transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
@@ -102,22 +115,41 @@ def build_dataloaders(
     use_clahe: bool,
     batch_size: int = 32,
     num_workers: int = 4,
+    img_size: int = 224,
 ):
-    """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))
+    """Create DataLoaders for all three splits.
+
+    The training DataLoader uses WeightedRandomSampler so that each class
+    is sampled at approximately equal frequency, counteracting class imbalance.
+    """
+    train_ds = DRDataset(train_paths, train_labels, transform=build_transforms(use_clahe, training=True, img_size=img_size))
+    val_ds = DRDataset(val_paths, val_labels, transform=build_transforms(use_clahe, training=False, img_size=img_size))
+    test_ds = DRDataset(test_paths, test_labels, transform=build_transforms(use_clahe, training=False, img_size=img_size))
+
+    # Per-sample weights: inverse of class frequency so all classes are seen equally
+    labels_arr = np.array(train_labels)
+    class_counts = np.bincount(labels_arr, minlength=5)
+    class_sample_weights = 1.0 / class_counts.astype(float)
+    sample_weights = class_sample_weights[labels_arr]
+    sampler = WeightedRandomSampler(
+        weights=torch.from_numpy(sample_weights).float(),
+        num_samples=len(train_labels),
+        replacement=True,
+    )
 
     train_loader = DataLoader(
-        train_ds, batch_size=batch_size, shuffle=True,
-        num_workers=num_workers, pin_memory=True
+        train_ds, batch_size=batch_size, sampler=sampler,
+        num_workers=num_workers, pin_memory=True,
+        persistent_workers=True, prefetch_factor=4,
     )
     val_loader = DataLoader(
         val_ds, batch_size=batch_size * 2, shuffle=False,
-        num_workers=num_workers, pin_memory=True
+        num_workers=num_workers, pin_memory=True,
+        persistent_workers=True, prefetch_factor=4,
     )
     test_loader = DataLoader(
         test_ds, batch_size=batch_size * 2, shuffle=False,
-        num_workers=num_workers, pin_memory=True
+        num_workers=num_workers, pin_memory=True,
+        persistent_workers=True, prefetch_factor=4,
     )
     return train_loader, val_loader, test_loader
diff --git a/evaluate.py b/evaluate.py
index 7a39db4..7f996dc 100644
--- a/evaluate.py
+++ b/evaluate.py
@@ -76,10 +76,10 @@ def plot_training_curves(
     ax1.set_title('Training and Validation Loss')
     ax1.legend()
 
-    ax2.plot(epochs, val_f1s, label='Val weighted F1', color='green')
+    ax2.plot(epochs, val_f1s, label='Val macro F1', color='green')
     ax2.set_xlabel('Epoch')
-    ax2.set_ylabel('Weighted F1')
-    ax2.set_title('Validation Weighted F1')
+    ax2.set_ylabel('Macro F1')
+    ax2.set_title('Validation Macro F1')
     ax2.legend()
 
     plt.tight_layout()
@@ -127,18 +127,28 @@ def generate_summary(all_results: dict, save_dir: str) -> None:
         writer.writerows(rows)
     print(f"\nSummary saved to {csv_path}")
 
-    # Comparison bar chart
-    labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows]
+    # Comparison bar chart — show both macro and weighted F1
+    x_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')
+    m_f1_vals = [r['macro_f1'] for r in rows]
+
+    x = np.arange(len(x_labels))
+    bar_width = 0.35
+    fig, ax = plt.subplots(figsize=(12, 5))
+    bars_w = ax.bar(x - bar_width / 2, w_f1_vals, bar_width, label='Weighted F1', color='#4C72B0')
+    bars_m = ax.bar(x + bar_width / 2, m_f1_vals, bar_width, label='Macro F1', color='#DD8452')
+    ax.set_ylabel('F1-Score')
+    ax.set_title('F1 Comparison Across Experiments')
+    ax.set_xticks(x)
+    ax.set_xticklabels(x_labels)
     ax.set_ylim(0, 1.0)
-    for bar, val in zip(bars, w_f1_vals):
+    ax.legend()
+    for bar, val in zip(bars_w, w_f1_vals):
+        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
+                f'{val:.3f}', ha='center', va='bottom', fontsize=8)
+    for bar, val in zip(bars_m, m_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)
+                f'{val:.3f}', ha='center', va='bottom', fontsize=8)
     plt.tight_layout()
     chart_path = os.path.join(save_dir, 'comparison_chart.png')
     fig.savefig(chart_path, dpi=150)
diff --git a/main.py b/main.py
index 6d8b2d9..fb87c33 100644
--- a/main.py
+++ b/main.py
@@ -12,10 +12,10 @@ import os
 
 import torch
 import torch.nn as nn
-from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
 import kagglehub
 
-from utils import set_seed, compute_class_weights, get_device, CLASS_NAMES
+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
@@ -35,15 +35,17 @@ def parse_args():
         '--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('--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=64, 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')
+    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')
     return parser.parse_args()
 
 
@@ -64,10 +66,6 @@ def main():
     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))
@@ -97,11 +95,35 @@ def main():
         )
 
         model = create_model(model_name, num_classes=5, pretrained=True).to(device)
-        criterion = nn.CrossEntropyLoss(weight=class_weights)
+
+        # 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(
-            model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+            [
+                {'params': backbone_params, 'lr': args.lr * 0.1},
+                {'params': head_params, 'lr': args.lr},
+            ],
+            weight_decay=args.weight_decay,
         )
-        scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs)
+        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(
@@ -113,16 +135,20 @@ def main():
             model_name=exp_key,
         )
 
-        # Save best checkpoint
+        # Save best checkpoint (strip DataParallel 'module.' prefix for portability)
         ckpt_path = os.path.join(exp_dir, 'best_model.pth')
-        torch.save(train_results['best_model_state'], ckpt_path)
+        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
-        model.load_state_dict(train_results['best_model_state'])
+        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(model, test_loader, device)
+        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} | "
diff --git a/train.py b/train.py
index 545b299..4a04a30 100644
--- a/train.py
+++ b/train.py
@@ -19,11 +19,11 @@ def train_model(
     patience: int = 5,
     model_name: str = "",
 ) -> dict:
-    """Train model with mixed precision, early stopping on val weighted-F1.
+    """Train model with mixed precision, early stopping on val macro F1.
 
     Returns a dict with:
-      best_model_state, train_losses, val_losses, val_f1s,
-      train_time (seconds), epochs_trained
+      best_model_state, train_losses, val_losses, val_f1s (macro), val_weighted_f1s,
+      train_time (seconds), epochs_trained, best_val_f1 (macro)
     """
     scaler = torch.amp.GradScaler('cuda') if device.type == 'cuda' else None
 
@@ -31,7 +31,7 @@ def train_model(
     best_state = None
     patience_counter = 0
 
-    train_losses, val_losses, val_f1s = [], [], []
+    train_losses, val_losses, val_f1s, val_weighted_f1s = [], [], [], []
     start_time = time.time()
 
     for epoch in range(1, num_epochs + 1):
@@ -97,21 +97,23 @@ def train_model(
         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)
+        val_f1_macro = f1_score(all_labels, all_preds, average='macro', zero_division=0)
+        val_f1_weighted = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
+        val_f1s.append(val_f1_macro)
+        val_weighted_f1s.append(val_f1_weighted)
 
         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}"
+            f"macro F1: {val_f1_macro:.4f} | weighted F1: {val_f1_weighted:.4f} | lr: {current_lr:.2e}"
         )
 
         scheduler.step()
 
-        # ── Early stopping ────────────────────────────────────────────────────
-        if val_f1 > best_val_f1:
-            best_val_f1 = val_f1
+        # ── Early stopping (monitored on macro F1) ───────────────────────────
+        if val_f1_macro > best_val_f1:
+            best_val_f1 = val_f1_macro
             best_state = copy.deepcopy(model.state_dict())
             patience_counter = 0
         else:
@@ -123,14 +125,15 @@ def train_model(
     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}")
+    print(f"  Training complete: {epochs_trained} epochs, {train_time:.1f}s, best val macro F1: {best_val_f1:.4f}")
 
     return {
         'best_model_state': best_state,
         'train_losses': train_losses,
         'val_losses': val_losses,
-        'val_f1s': val_f1s,
+        'val_f1s': val_f1s,           # macro F1 per epoch (used for early stopping)
+        'val_weighted_f1s': val_weighted_f1s,
         'train_time': train_time,
         'epochs_trained': epochs_trained,
-        'best_val_f1': best_val_f1,
+        'best_val_f1': best_val_f1,   # best macro F1
     }
diff --git a/utils.py b/utils.py
index a836ca2..90c79d5 100644
--- a/utils.py
+++ b/utils.py
@@ -1,6 +1,8 @@
 import random
 import numpy as np
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 
 CLASS_NAMES = [
     'Healthy (0)',
@@ -18,8 +20,7 @@ def set_seed(seed: int = 42) -> None:
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.benchmark = True
 
 
 def compute_class_weights(labels) -> torch.FloatTensor:
@@ -35,3 +36,28 @@ def get_device() -> torch.device:
     if torch.cuda.is_available():
         return torch.device('cuda')
     return torch.device('cpu')
+
+
+class FocalLoss(nn.Module):
+    """Focal loss with optional per-class alpha weighting.
+
+    Focuses learning on hard, misclassified examples by down-weighting
+    easy examples (high pt). Good for class imbalance in medical imaging.
+
+    Args:
+        alpha: Per-class weight tensor (same shape as class weights for CE).
+               If None, no per-class weighting is applied.
+        gamma: Focusing parameter. gamma=0 reduces to standard CE.
+               gamma=2 is the standard value from the RetinaNet paper.
+    """
+
+    def __init__(self, alpha=None, gamma: float = 2.0):
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        ce_loss = F.cross_entropy(inputs, targets, weight=self.alpha, reduction='none')
+        pt = torch.exp(-ce_loss)
+        focal_loss = ((1.0 - pt) ** self.gamma) * ce_loss
+        return focal_loss.mean()