File size: 16,562 Bytes
fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 48572da fac3f86 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 | """
Shared dataset classes and loading utilities for GAP-CLIP evaluation scripts.
Provides:
- FashionMNISTDataset (Fashion-MNIST grayscale images)
- KaggleDataset (KAGL Marqo HuggingFace dataset)
- LocalDataset (internal local validation dataset)
- Matching load_* convenience functions
- collate_fn_filter_none (for DataLoader)
- normalize_hierarchy_label (text normalisation helper)
"""
from __future__ import annotations
import difflib
import hashlib
import os
import sys
from pathlib import Path
from io import BytesIO
from typing import List, Optional
import numpy as np
import pandas as pd
import torch
from PIL import Image
import requests
from torch.utils.data import Dataset
from torchvision import transforms
# Make project root importable when running evaluation scripts directly.
_PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(_PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(_PROJECT_ROOT))
from config import ( # type: ignore
ROOT_DIR,
column_local_image_path,
fashion_mnist_csv,
local_dataset_path,
images_dir,
)
# ---------------------------------------------------------------------------
# Fashion-MNIST helpers
# ---------------------------------------------------------------------------
def get_fashion_mnist_labels() -> dict:
"""Return the 10 Fashion-MNIST integer-to-name mapping."""
return {
0: "T-shirt/top",
1: "Trouser",
2: "Pullover",
3: "Dress",
4: "Coat",
5: "Sandal",
6: "Shirt",
7: "Sneaker",
8: "Bag",
9: "Ankle boot",
}
def create_fashion_mnist_to_hierarchy_mapping(hierarchy_classes: List[str]) -> dict:
"""Map Fashion-MNIST integer labels to nearest hierarchy class name.
Returns dict {label_id: matched_class_name or None}.
"""
fashion_mnist_labels = get_fashion_mnist_labels()
hierarchy_classes_lower = [h.lower() for h in hierarchy_classes]
mapping = {}
for fm_label_id, fm_label in fashion_mnist_labels.items():
fm_label_lower = fm_label.lower()
matched_hierarchy = None
if fm_label_lower in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index(fm_label_lower)]
elif any(h in fm_label_lower or fm_label_lower in h for h in hierarchy_classes_lower):
for h_class in hierarchy_classes:
if h_class.lower() in fm_label_lower or fm_label_lower in h_class.lower():
matched_hierarchy = h_class
break
else:
if fm_label_lower in ["t-shirt/top", "top"]:
if "top" in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index("top")]
elif "trouser" in fm_label_lower:
for p in ["bottom", "pants", "trousers", "trouser", "pant"]:
if p in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index(p)]
break
elif "pullover" in fm_label_lower:
for p in ["sweater", "pullover"]:
if p in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index(p)]
break
elif "dress" in fm_label_lower:
if "dress" in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index("dress")]
elif "coat" in fm_label_lower:
for p in ["jacket", "outerwear", "coat"]:
if p in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index(p)]
break
elif fm_label_lower in ["sandal", "sneaker", "ankle boot"]:
for p in ["shoes", "shoe", "sandal", "sneaker", "boot"]:
if p in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index(p)]
break
elif "bag" in fm_label_lower:
if "bag" in hierarchy_classes_lower:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index("bag")]
if matched_hierarchy is None:
close = difflib.get_close_matches(fm_label_lower, hierarchy_classes_lower, n=1, cutoff=0.6)
if close:
matched_hierarchy = hierarchy_classes[hierarchy_classes_lower.index(close[0])]
mapping[fm_label_id] = matched_hierarchy
status = matched_hierarchy if matched_hierarchy else "NO MATCH (will be filtered out)"
print(f" {fm_label} ({fm_label_id}) -> {status}")
return mapping
def convert_fashion_mnist_to_image(pixel_values) -> Image.Image:
"""Convert a flat 784-element pixel array to an RGB PIL image."""
arr = np.array(pixel_values).reshape(28, 28).astype(np.uint8)
arr = np.stack([arr] * 3, axis=-1)
return Image.fromarray(arr)
class FashionMNISTDataset(Dataset):
"""PyTorch dataset wrapping Fashion-MNIST CSV rows."""
def __init__(self, dataframe: pd.DataFrame, image_size: int = 224, label_mapping: Optional[dict] = None):
self.dataframe = dataframe
self.image_size = image_size
self.labels_map = get_fashion_mnist_labels()
self.label_mapping = label_mapping
self.transform = 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]),
])
def __len__(self) -> int:
return len(self.dataframe)
def __getitem__(self, idx):
row = self.dataframe.iloc[idx]
pixel_cols = [f"pixel{i}" for i in range(1, 785)]
image = convert_fashion_mnist_to_image(row[pixel_cols].values)
image = self.transform(image)
label_id = int(row["label"])
description = self.labels_map[label_id]
color = "unknown"
hierarchy = (
self.label_mapping[label_id]
if (self.label_mapping and label_id in self.label_mapping)
else self.labels_map[label_id]
)
return image, description, color, hierarchy
def load_fashion_mnist_dataset(
max_samples: int = 10000,
hierarchy_classes: Optional[List[str]] = None,
csv_path: Optional[str] = None,
) -> FashionMNISTDataset:
"""Load Fashion-MNIST test CSV into a FashionMNISTDataset.
Args:
max_samples: Maximum number of samples to use.
hierarchy_classes: If provided, maps Fashion-MNIST labels to these classes.
csv_path: Path to fashion-mnist_test.csv. Defaults to config.fashion_mnist_csv.
"""
if csv_path is None:
csv_path = fashion_mnist_csv
print("Loading Fashion-MNIST test dataset...")
df = pd.read_csv(csv_path)
print(f"Fashion-MNIST dataset loaded: {len(df)} samples")
label_mapping = None
if hierarchy_classes is not None:
print("\nCreating mapping from Fashion-MNIST labels to hierarchy classes:")
label_mapping = create_fashion_mnist_to_hierarchy_mapping(hierarchy_classes)
valid_ids = [lid for lid, h in label_mapping.items() if h is not None]
df = df[df["label"].isin(valid_ids)]
print(f"\nAfter filtering to mappable labels: {len(df)} samples")
df_sample = df.head(max_samples)
print(f"Using {len(df_sample)} samples for evaluation")
return FashionMNISTDataset(df_sample, label_mapping=label_mapping)
# ---------------------------------------------------------------------------
# KAGL Marqo dataset
# ---------------------------------------------------------------------------
class KaggleDataset(Dataset):
"""Dataset class for KAGL Marqo HuggingFace dataset."""
def __init__(self, dataframe: pd.DataFrame, image_size: int = 224, include_hierarchy: bool = False):
self.dataframe = dataframe
self.image_size = image_size
self.include_hierarchy = include_hierarchy
self.transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def __len__(self) -> int:
return len(self.dataframe)
def __getitem__(self, idx):
row = self.dataframe.iloc[idx]
image_data = row["image_url"]
if isinstance(image_data, dict) and "bytes" in image_data:
image = Image.open(BytesIO(image_data["bytes"])).convert("RGB")
elif hasattr(image_data, "convert"):
image = image_data.convert("RGB")
else:
image = Image.open(BytesIO(image_data)).convert("RGB")
image = self.transform(image)
description = row["text"]
color = row["color"]
if self.include_hierarchy:
hierarchy = row.get("hierarchy", "unknown")
return image, description, color, hierarchy
return image, description, color
def download_kaggle_raw_df() -> pd.DataFrame:
"""Download the raw KAGL Marqo DataFrame from HuggingFace.
This is the expensive network operation. Callers can cache the result
and pass it to :func:`load_kaggle_marqo_dataset` via *raw_df* to avoid
repeated downloads.
"""
from datasets import load_dataset # type: ignore
print("Downloading KAGL Marqo dataset from HuggingFace...")
dataset = load_dataset("Marqo/KAGL")
df = dataset["data"].to_pandas()
print(f"KAGL dataset downloaded: {len(df)} samples, columns: {list(df.columns)}")
return df
def load_kaggle_marqo_dataset(
max_samples: int = 5000,
include_hierarchy: bool = False,
raw_df: Optional[pd.DataFrame] = None,
) -> KaggleDataset:
"""Download and prepare the KAGL Marqo HuggingFace dataset.
Args:
max_samples: Maximum number of samples to return.
include_hierarchy: If True, dataset tuples include a hierarchy element.
raw_df: Pre-downloaded DataFrame (from :func:`download_kaggle_raw_df`).
If *None*, the dataset is downloaded from HuggingFace.
"""
if raw_df is not None:
df = raw_df.copy()
print(f"Using cached KAGL DataFrame: {len(df)} samples")
else:
df = download_kaggle_raw_df()
df = df.dropna(subset=["text", "image"])
if len(df) > max_samples:
df = df.sample(n=max_samples, random_state=42)
print(f"Sampled {max_samples} items")
kaggle_df = pd.DataFrame({
"image_url": df["image"],
"text": df["text"],
"color": df["baseColour"].str.lower().str.replace("grey", "gray"),
})
kaggle_df = kaggle_df.dropna(subset=["color"])
print(f"Colors: {sorted(kaggle_df['color'].unique())}")
return KaggleDataset(kaggle_df, include_hierarchy=include_hierarchy)
# ---------------------------------------------------------------------------
# Local validation dataset
# ---------------------------------------------------------------------------
class LocalDataset(Dataset):
"""Dataset class for the internal local validation dataset."""
def __init__(self, dataframe: pd.DataFrame, image_size: int = 224, include_hierarchy: bool = False):
self.dataframe = dataframe
self.image_size = image_size
self.include_hierarchy = include_hierarchy
self.transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def __len__(self) -> int:
return len(self.dataframe)
def __getitem__(self, idx):
row = self.dataframe.iloc[idx]
try:
image_path = row.get(column_local_image_path) if hasattr(row, "get") else None
if isinstance(image_path, str) and image_path:
if not os.path.isabs(image_path):
image_path = str(ROOT_DIR / image_path)
image = Image.open(image_path).convert("RGB")
else:
# Fallback: download image from URL (and cache).
image_url = row.get("image_url") if hasattr(row, "get") else None
if isinstance(image_url, dict) and "bytes" in image_url:
image = Image.open(BytesIO(image_url["bytes"])).convert("RGB")
elif isinstance(image_url, str) and image_url:
cache_dir = Path(images_dir)
cache_dir.mkdir(parents=True, exist_ok=True)
url_hash = hashlib.md5(image_url.encode("utf-8")).hexdigest()
cache_path = cache_dir / f"{url_hash}.jpg"
if cache_path.exists():
image = Image.open(cache_path).convert("RGB")
else:
resp = requests.get(image_url, timeout=10)
resp.raise_for_status()
image = Image.open(BytesIO(resp.content)).convert("RGB")
image.save(cache_path, "JPEG", quality=85, optimize=True)
else:
raise ValueError("Missing image_path and image_url")
except Exception as e:
print(f"Error loading image: {e}")
image = Image.new("RGB", (224, 224), color="gray")
image = self.transform(image)
description = row["text"]
color = row["color"]
if self.include_hierarchy:
hierarchy = row.get("hierarchy", "unknown")
return image, description, color, hierarchy
return image, description, color
def load_local_validation_dataset(
max_samples: int = 5000,
include_hierarchy: bool = False,
raw_df: Optional[pd.DataFrame] = None,
) -> LocalDataset:
"""Load and prepare the internal local validation dataset.
Args:
max_samples: Maximum number of samples to return.
include_hierarchy: If True, dataset tuples include a hierarchy element.
raw_df: Pre-loaded DataFrame. If *None*, the CSV is read from disk.
"""
if raw_df is not None:
df = raw_df.copy()
print(f"Using cached local DataFrame: {len(df)} samples")
else:
print("Loading local validation dataset...")
df = pd.read_csv(local_dataset_path)
print(f"Dataset loaded: {len(df)} samples")
if column_local_image_path in df.columns:
df = df.dropna(subset=[column_local_image_path])
print(f"After filtering NaN image paths: {len(df)} samples")
else:
print(f"Column '{column_local_image_path}' not found; falling back to 'image_url'.")
if "color" in df.columns:
print(f"After color filtering: {len(df)} samples, colors: {sorted(df['color'].unique())}")
if len(df) > max_samples:
df = df.sample(n=max_samples, random_state=42)
print(f"Sampled {max_samples} items")
print(f"Using {len(df)} samples for evaluation")
return LocalDataset(df, include_hierarchy=include_hierarchy)
# ---------------------------------------------------------------------------
# DataLoader utilities
# ---------------------------------------------------------------------------
def collate_fn_filter_none(batch):
"""Collate function that silently drops None items from a batch."""
original_len = len(batch)
batch = [item for item in batch if item is not None]
if original_len > len(batch):
print(f"Filtered out {original_len - len(batch)} None values from batch")
if not batch:
print("Empty batch after filtering None values")
return torch.tensor([]), [], []
# Support both 3-value (image, text, color) and 4-value (image, text, color, hierarchy) tuples
if len(batch[0]) == 4:
images, texts, colors, hierarchies = zip(*batch)
return torch.stack(images), list(texts), list(colors), list(hierarchies)
images, texts, colors = zip(*batch)
return torch.stack(images), list(texts), list(colors)
# ---------------------------------------------------------------------------
# Text normalisation helpers
# ---------------------------------------------------------------------------
def normalize_hierarchy_label(label: str) -> str:
"""Lower-case and strip a hierarchy label for consistent comparison."""
return label.lower().strip() if label else ""
|