1 files changed, 51 insertions, 19 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