Update data_utils.py

data_utils.py

@@ -1,16 +1,20 @@
-from typing import List, Tuple
+import random
+from collections import Counter
+from typing import Dict, List, Tuple
 
+import pandas as pd
 import torch
 from PIL import Image
-from datasets import load_dataset
-from torch.utils.data import Dataset, DataLoader, random_split
+from datasets import load_dataset, DatasetDict
+from torch.utils.data import Dataset, DataLoader, Subset
 from torchvision import transforms
 
 from config import HF_DATASET_REPO, HF_TOKEN, IMAGE_SIZE, RANDOM_SEED
 
 
+_RAW_DATASET = None
 _CLASS_NAMES = None
-_HF_DATASET_CACHE = None
+_SPLITS = None
 
 
 class HFDatasetWrapper(Dataset):
@@ -31,13 +35,24 @@ class HFDatasetWrapper(Dataset):
         image = image.convert("RGB")
         label = int(item["label"])
 
-        return self.transform(image), label
+        if self.transform:
+            image = self.transform(image)
 
+        return image, label
 
-def get_transform():
+
+def get_train_transform():
     return transforms.Compose(
         [
-            transforms.Resize((224, 224)),
+            transforms.RandomResizedCrop(IMAGE_SIZE, scale=(0.65, 1.0)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.RandomVerticalFlip(p=0.2),
+            transforms.RandomRotation(degrees=15),
+            transforms.ColorJitter(
+                brightness=0.2,
+                contrast=0.2,
+                saturation=0.1,
+            ),
             transforms.ToTensor(),
             transforms.Normalize(
                 mean=(0.485, 0.456, 0.406),
@@ -47,70 +62,219 @@ def get_transform():
     )
 
 
-def load_charcoal_dataset():
-    global _CLASS_NAMES, _HF_DATASET_CACHE
+def get_eval_transform():
+    return transforms.Compose(
+        [
+            transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=(0.485, 0.456, 0.406),
+                std=(0.229, 0.224, 0.225),
+            ),
+        ]
+    )
+
 
-    if _HF_DATASET_CACHE is not None:
-        return _HF_DATASET_CACHE, _CLASS_NAMES
+def load_raw_dataset():
+    global _RAW_DATASET, _CLASS_NAMES
+
+    if _RAW_DATASET is not None:
+        return _RAW_DATASET, _CLASS_NAMES
 
     if not HF_TOKEN:
         raise RuntimeError(
-            "HF_TOKEN is missing. Please add it in the Space secrets."
+            "HF_TOKEN est manquant. Ajoutez-le dans les Secrets du Space Hugging Face."
         )
 
     raw = load_dataset(HF_DATASET_REPO, token=HF_TOKEN)
 
+    if "train" not in raw:
+        raise RuntimeError("Le dataset Hugging Face doit contenir au moins un split 'train'.")
+
     label_feature = raw["train"].features["label"]
-    if hasattr(label_feature, "names"):
-        _CLASS_NAMES = label_feature.names
+
+    if hasattr(label_feature, "names") and label_feature.names:
+        class_names = label_feature.names
     else:
-        _CLASS_NAMES = sorted(list(set(raw["train"]["label"])))
+        labels = list(raw["train"]["label"])
+        class_names = [str(x) for x in sorted(set(labels))]
+
+    _RAW_DATASET = raw
+    _CLASS_NAMES = class_names
+
+    return _RAW_DATASET, _CLASS_NAMES
+
+
+def prepare_splits(
+    train_ratio: float = 0.80,
+    val_ratio: float = 0.10,
+    test_ratio: float = 0.10,
+):
+    global _SPLITS
+
+    if _SPLITS is not None:
+        return _SPLITS
+
+    raw, class_names = load_raw_dataset()
+
+    if "validation" in raw and "test" in raw:
+        _SPLITS = {
+            "train": raw["train"],
+            "validation": raw["validation"],
+            "test": raw["test"],
+        }
+        return _SPLITS
+
+    if "test" in raw:
+        train_val = raw["train"]
+        test = raw["test"]
+
+        relative_val_ratio = val_ratio / (train_ratio + val_ratio)
 
-    if "test" not in raw:
         try:
-            split = raw["train"].train_test_split(
-                test_size=0.2,
+            split_train_val = train_val.train_test_split(
+                test_size=relative_val_ratio,
                 seed=RANDOM_SEED,
                 stratify_by_column="label",
             )
         except Exception:
-            split = raw["train"].train_test_split(
-                test_size=0.2,
+            split_train_val = train_val.train_test_split(
+                test_size=relative_val_ratio,
                 seed=RANDOM_SEED,
             )
 
-        raw = {
-            "train": split["train"],
-            "test": split["test"],
+        _SPLITS = {
+            "train": split_train_val["train"],
+            "validation": split_train_val["test"],
+            "test": test,
         }
+        return _SPLITS
+
+    full = raw["train"]
+
+    try:
+        first_split = full.train_test_split(
+            test_size=(val_ratio + test_ratio),
+            seed=RANDOM_SEED,
+            stratify_by_column="label",
+        )
+    except Exception:
+        first_split = full.train_test_split(
+            test_size=(val_ratio + test_ratio),
+            seed=RANDOM_SEED,
+        )
+
+    temp = first_split["test"]
+    relative_test_ratio = test_ratio / (val_ratio + test_ratio)
 
-    _HF_DATASET_CACHE = raw
-    return _HF_DATASET_CACHE, _CLASS_NAMES
+    try:
+        second_split = temp.train_test_split(
+            test_size=relative_test_ratio,
+            seed=RANDOM_SEED,
+            stratify_by_column="label",
+        )
+    except Exception:
+        second_split = temp.train_test_split(
+            test_size=relative_test_ratio,
+            seed=RANDOM_SEED,
+        )
+
+    _SPLITS = {
+        "train": first_split["train"],
+        "validation": second_split["train"],
+        "test": second_split["test"],
+    }
+
+    return _SPLITS
 
 
 def get_class_names() -> List[str]:
-    _, class_names = load_charcoal_dataset()
+    _, class_names = load_raw_dataset()
     return class_names
 
 
-def make_loaders(batch_size: int, val_ratio: float = 0.1):
-    raw, class_names = load_charcoal_dataset()
-    transform = get_transform()
+def make_loaders(batch_size: int):
+    splits = prepare_splits()
+    class_names = get_class_names()
 
-    train_dataset = HFDatasetWrapper(raw["train"], transform)
-    test_dataset = HFDatasetWrapper(raw["test"], transform)
+    train_dataset = HFDatasetWrapper(splits["train"], get_train_transform())
+    val_dataset = HFDatasetWrapper(splits["validation"], get_eval_transform())
+    test_dataset = HFDatasetWrapper(splits["test"], get_eval_transform())
 
-    val_size = int(len(train_dataset) * val_ratio)
-    train_size = len(train_dataset) - val_size
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
-    train_subset, val_subset = random_split(
-        train_dataset,
-        [train_size, val_size],
-        generator=torch.Generator().manual_seed(RANDOM_SEED),
-    )
+    return train_loader, val_loader, test_loader, class_names
 
-    train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False)
-    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
 
-    return train_loader, val_loader, test_loader, class_names
+def dataset_overview() -> Tuple[dict, pd.DataFrame]:
+    splits = prepare_splits()
+    class_names = get_class_names()
+
+    rows = []
+    total = 0
+
+    for split_name, split_data in splits.items():
+        labels = list(split_data["label"])
+        counter = Counter(labels)
+        split_total = len(labels)
+        total += split_total
+
+        for label_id, count in sorted(counter.items()):
+            rows.append(
+                {
+                    "split": split_name,
+                    "classe": class_names[int(label_id)],
+                    "nombre_images": count,
+                }
+            )
+
+    df = pd.DataFrame(rows)
+
+    summary = {
+        "dataset": HF_DATASET_REPO,
+        "nombre_total_images": total,
+        "nombre_classes": len(class_names),
+        "train": len(splits["train"]),
+        "validation": len(splits["validation"]),
+        "test": len(splits["test"]),
+    }
+
+    return summary, df
+
+
+def get_images_for_gallery(split_name: str, class_name: str, max_images: int = 24):
+    splits = prepare_splits()
+    class_names = get_class_names()
+
+    if split_name not in splits:
+        split_name = "train"
+
+    dataset = splits[split_name]
+
+    if class_name and class_name != "Toutes les classes":
+        class_id = class_names.index(class_name)
+        indices = [i for i, x in enumerate(dataset["label"]) if int(x) == class_id]
+    else:
+        indices = list(range(len(dataset)))
+
+    if not indices:
+        return []
+
+    sample_indices = random.sample(indices, min(max_images, len(indices)))
+
+    gallery = []
+    for idx in sample_indices:
+        item = dataset[idx]
+        image = item["image"]
+        if not isinstance(image, Image.Image):
+            image = Image.open(image)
+        image = image.convert("RGB")
+
+        label_id = int(item["label"])
+        label_name = class_names[label_id]
+
+        gallery.append((image, label_name))
+
+    return gallery