Spaces:
Running
Running
| from pathlib import Path | |
| from typing import List, Tuple | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from sklearn.model_selection import train_test_split | |
| from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler | |
| class FishDataset(Dataset): | |
| """ | |
| Accepts two folder layouts: | |
| Auto-split (stratified 80/20 by default): | |
| data/raw/<class_name>/*.jpg | |
| Pre-split: | |
| data/raw/train/<class_name>/*.jpg | |
| data/raw/val/<class_name>/*.jpg | |
| """ | |
| IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} | |
| def __init__( | |
| self, | |
| root_dir: str, | |
| transform=None, | |
| split: str = "train", | |
| val_ratio: float = 0.2, | |
| seed: int = 42, | |
| ): | |
| self.transform = transform | |
| root = Path(root_dir) | |
| if (root / "train").exists() and (root / "val").exists(): | |
| data_dir = root / split | |
| self.classes = sorted(d.name for d in data_dir.iterdir() if d.is_dir()) | |
| self.class_to_idx = {c: i for i, c in enumerate(self.classes)} | |
| self.samples = self._scan(data_dir) | |
| else: | |
| self.classes = sorted(d.name for d in root.iterdir() if d.is_dir()) | |
| self.class_to_idx = {c: i for i, c in enumerate(self.classes)} | |
| all_samples = self._scan(root) | |
| labels = [s[1] for s in all_samples] | |
| train_s, val_s = train_test_split( | |
| all_samples, test_size=val_ratio, random_state=seed, stratify=labels | |
| ) | |
| self.samples = train_s if split == "train" else val_s | |
| def _scan(self, directory: Path) -> List[Tuple[Path, int]]: | |
| samples = [] | |
| for cls_dir in sorted(directory.iterdir()): | |
| if not cls_dir.is_dir() or cls_dir.name not in self.class_to_idx: | |
| continue | |
| idx = self.class_to_idx[cls_dir.name] | |
| for p in cls_dir.iterdir(): | |
| if p.suffix.lower() in self.IMG_EXTS: | |
| samples.append((p, idx)) | |
| return samples | |
| def __len__(self) -> int: | |
| return len(self.samples) | |
| def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]: | |
| path, label = self.samples[idx] | |
| image = np.array(Image.open(path).convert("RGB")) | |
| if self.transform: | |
| image = self.transform(image=image)["image"] | |
| return image, label | |
| def sample_weights(self) -> torch.Tensor: | |
| counts = np.zeros(len(self.classes)) | |
| for _, lbl in self.samples: | |
| counts[lbl] += 1 | |
| w = 1.0 / np.maximum(counts, 1) | |
| return torch.tensor([w[lbl] for _, lbl in self.samples], dtype=torch.float32) | |
| def build_dataloaders(config: dict) -> Tuple[DataLoader, DataLoader, List[str]]: | |
| from src.transforms import get_train_transforms, get_val_transforms | |
| img_size = config["data"]["image_size"] | |
| val_ratio = 1.0 - config["data"]["train_split"] | |
| bs = config["training"]["batch_size"] | |
| nw = config["data"]["num_workers"] | |
| train_ds = FishDataset(config["data"]["data_dir"], get_train_transforms(img_size), "train", val_ratio) | |
| val_ds = FishDataset(config["data"]["data_dir"], get_val_transforms(img_size), "val", val_ratio) | |
| sampler = WeightedRandomSampler(train_ds.sample_weights(), len(train_ds)) | |
| train_loader = DataLoader(train_ds, batch_size=bs, sampler=sampler, num_workers=nw, pin_memory=True) | |
| val_loader = DataLoader(val_ds, batch_size=bs, shuffle=False, num_workers=nw, pin_memory=True) | |
| return train_loader, val_loader, train_ds.classes | |