Droid210/FleetVision

1"""Data loading and augmentation utilities."""2from pathlib import Path3from typing import List, Tuple4 5import torch6from torch.utils.data import DataLoader, Dataset7from torchvision import transforms8from transformers import AutoImageProcessor9 10from .config import EXPECTED_CLASSES, IMAGE_SIZE, MODEL_ID11 12 13class DamagedImageAugmentation:14    """Heavy augmentation for Damaged class to simulate different angles."""15 16    def __call__(self, img):17        """Apply random augmentations."""18        img = transforms.RandomResizedCrop(IMAGE_SIZE, scale=(0.7, 1.0))(img)19        img = transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1)(img)20        img = transforms.RandomRotation(25)(img)21        img = transforms.RandomAffine(degrees=0, translate=(0.1, 0.1))(img)22        img = transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 0.5))(img)23        return img24 25 26class DamageDataset(Dataset):27    """Custom dataset for damage detection with class-specific augmentation."""28 29    def __init__(30        self,31        data_dir: Path,32        processor,33        is_train: bool = True,34    ):35        """Initialize dataset.36 37        Args:38            data_dir: Path to dataset root (parent of Whole/Damaged folders).39            processor: Hugging Face image processor for preprocessing.40            is_train: Whether this is training set.41        """42        self.processor = processor43        self.is_train = is_train44        self.images = []45        self.labels = []46 47        # Load images and labels48        for class_idx, class_name in enumerate(EXPECTED_CLASSES):49            class_dir = data_dir / class_name50            if not class_dir.exists():51                raise FileNotFoundError(f"Class directory not found: {class_dir}")52 53            for img_path in class_dir.glob("*.jpg"):54                self.images.append(img_path)55                self.labels.append(class_idx)56 57            for img_path in class_dir.glob("*.png"):58                self.images.append(img_path)59                self.labels.append(class_idx)60 61        if not self.images:62            raise ValueError(f"No images found in {data_dir}")63 64        # Drop stale paths if files were moved/deleted after listing.65        filtered = [(p, y) for p, y in zip(self.images, self.labels) if p.exists() and p.is_file()]66        self.images = [p for p, _ in filtered]67        self.labels = [y for _, y in filtered]68 69        if not self.images:70            raise ValueError(f"No readable image files found in {data_dir}")71 72        # Augmentation for training73        self.damaged_aug = DamagedImageAugmentation() if is_train else None74 75    def __len__(self) -> int:76        return len(self.images)77 78    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:79        """Get image and label.80 81        Args:82            idx: Index of image.83 84        Returns:85            Tuple of (image_tensor, label_tensor).86        """87        from PIL import Image88 89        max_tries = min(16, len(self.images))90        for offset in range(max_tries):91            safe_idx = (idx + offset) % len(self.images)92            img_path = self.images[safe_idx]93            label = self.labels[safe_idx]94 95            if not img_path.exists():96                continue97 98            try:99                image = Image.open(img_path).convert("RGB")100            except (FileNotFoundError, OSError):101                continue102 103            # Apply heavy augmentation to Damaged class during training104            if self.is_train and label == 1:  # Damaged class105                image = self.damaged_aug(image)106 107            # Process with ViTImageProcessor108            processed = self.processor(image, return_tensors="pt")109            pixel_values = processed["pixel_values"].squeeze(0)110            return pixel_values, torch.tensor(label, dtype=torch.long)111 112        raise FileNotFoundError("Failed to load a valid image sample after multiple attempts.")113 114 115class BinaryDamageDataset(Dataset):116    """Binary dataset built from explicit (image_path, label) pairs."""117 118    def __init__(119        self,120        samples: List[Tuple[Path, int, bool]],121        processor,122        is_train: bool = True,123    ):124        self.samples = [(p, y, syn) for p, y, syn in samples if p.exists() and p.is_file()]125        self.processor = processor126        self.is_train = is_train127        self.damaged_aug = DamagedImageAugmentation() if is_train else None128        self.synthetic_negative_aug = transforms.Compose(129            [130                transforms.RandomResizedCrop(IMAGE_SIZE, scale=(0.2, 0.6)),131                transforms.ColorJitter(brightness=0.6, contrast=0.6, saturation=0.6, hue=0.2),132                transforms.RandomGrayscale(p=0.4),133                transforms.GaussianBlur(kernel_size=5, sigma=(0.3, 1.5)),134            ]135        )136 137        if not self.samples:138            raise ValueError("No valid image files found after filtering sample paths.")139 140    def __len__(self) -> int:141        return len(self.samples)142 143    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:144        from PIL import Image145 146        max_tries = min(16, len(self.samples))147        for offset in range(max_tries):148            safe_idx = (idx + offset) % len(self.samples)149            img_path, label, is_synthetic_negative = self.samples[safe_idx]150 151            if not img_path.exists():152                continue153 154            try:155                image = Image.open(img_path).convert("RGB")156            except (FileNotFoundError, OSError):157                continue158 159            if self.is_train and label == 1:  # Damaged class160                image = self.damaged_aug(image)161            elif is_synthetic_negative:162                # Damaged-only fallback: synthesize non-damage-like negatives.163                image = self.synthetic_negative_aug(image)164 165            processed = self.processor(image, return_tensors="pt")166            pixel_values = processed["pixel_values"].squeeze(0)167            return pixel_values, torch.tensor(label, dtype=torch.long)168 169        raise FileNotFoundError("Failed to load a valid image sample after multiple attempts.")170 171 172def _collect_images(folder: Path) -> List[Path]:173    """Collect images from a folder recursively."""174    exts = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp")175    files: List[Path] = []176    for ext in exts:177        files.extend(folder.rglob(ext))178    return sorted(files)179 180 181def _split_items(items: List[Path], val_ratio: float = 0.2, seed: int = 42) -> Tuple[List[Path], List[Path]]:182    """Deterministically split items into train/validation."""183    if len(items) < 2:184        return items, []185 186    g = torch.Generator().manual_seed(seed)187    perm = torch.randperm(len(items), generator=g).tolist()188    split_at = max(1, int(len(items) * (1.0 - val_ratio)))189    train_idx = perm[:split_at]190    val_idx = perm[split_at:]191 192    train_items = [items[i] for i in train_idx]193    val_items = [items[i] for i in val_idx]194    return train_items, val_items195 196 197def _resolve_damage_assessment_root(data_dir: Path) -> Path | None:198    """Resolve CarDD/FiftyOne-style dataset root."""199    if (data_dir / "samples.json").exists() and (data_dir / "data").exists():200        return data_dir201 202    nested = data_dir / "damage_assessment"203    if (nested / "samples.json").exists() and (nested / "data").exists():204        return nested205 206    return None207 208 209def _build_from_damage_assessment(210    data_dir: Path,211    batch_size: int,212    num_workers: int,213    processor,214) -> Tuple[DataLoader, DataLoader, object]:215    """Build loaders from CarDD/FiftyOne-style layout.216 217    Expected positives: <root>/damage_assessment/data or <root>/data218    Expected negatives: any of219      - <data_dir>/whole_pool220      - <data_dir>/Whole221      - <data_dir>/train/Whole + <data_dir>/valid/Whole222    """223    assessment_root = _resolve_damage_assessment_root(data_dir)224    if assessment_root is None:225        raise ValueError("Could not resolve damage_assessment dataset root.")226 227    damaged_images = _collect_images(assessment_root / "data")228    if not damaged_images:229        raise ValueError(f"No images found in {assessment_root / 'data'}")230 231    whole_candidates = [232        data_dir / "whole_pool",233        data_dir / "Whole",234        assessment_root.parent / "whole_pool",235        assessment_root.parent / "Whole",236    ]237 238    whole_images: List[Path] = []239    for folder in whole_candidates:240        if folder.exists():241            whole_images.extend(_collect_images(folder))242 243    train_whole = data_dir / "train" / "Whole"244    valid_whole = data_dir / "valid" / "Whole"245    parent_train_whole = assessment_root.parent / "train" / "Whole"246    parent_valid_whole = assessment_root.parent / "valid" / "Whole"247    if train_whole.exists():248        whole_images.extend(_collect_images(train_whole))249    if valid_whole.exists():250        whole_images.extend(_collect_images(valid_whole))251    if parent_train_whole.exists():252        whole_images.extend(_collect_images(parent_train_whole))253    if parent_valid_whole.exists():254        whole_images.extend(_collect_images(parent_valid_whole))255 256    whole_images = sorted(set(whole_images))257    train_damaged, val_damaged = _split_items(damaged_images, val_ratio=0.2, seed=42)258    train_samples: List[Tuple[Path, int, bool]]259    val_samples: List[Tuple[Path, int, bool]]260 261    if whole_images:262        train_whole_split, val_whole_split = _split_items(whole_images, val_ratio=0.2, seed=42)263        train_samples = [(p, 1, False) for p in train_damaged] + [(p, 0, False) for p in train_whole_split]264        val_samples = [(p, 1, False) for p in val_damaged] + [(p, 0, False) for p in val_whole_split]265    else:266        print(267            "WARNING: No Whole images found. Using damaged-only fallback with synthetic Whole negatives."268        )269        # Each damaged image is used twice: once as Damaged, once as synthetic Whole.270        train_samples = [(p, 1, False) for p in train_damaged] + [(p, 0, True) for p in train_damaged]271        val_samples = [(p, 1, False) for p in val_damaged] + [(p, 0, True) for p in val_damaged]272 273    if not train_samples or not val_samples:274        raise ValueError("Insufficient samples after split. Need at least one train and one validation sample per run.")275 276    g = torch.Generator().manual_seed(42)277    train_perm = torch.randperm(len(train_samples), generator=g).tolist()278    val_perm = torch.randperm(len(val_samples), generator=g).tolist()279    train_samples = [train_samples[i] for i in train_perm]280    val_samples = [val_samples[i] for i in val_perm]281 282    train_dataset = BinaryDamageDataset(train_samples, processor, is_train=True)283    val_dataset = BinaryDamageDataset(val_samples, processor, is_train=False)284 285    train_loader = DataLoader(286        train_dataset,287        batch_size=batch_size,288        shuffle=True,289        num_workers=num_workers,290        pin_memory=torch.cuda.is_available(),291    )292 293    val_loader = DataLoader(294        val_dataset,295        batch_size=batch_size,296        shuffle=False,297        num_workers=num_workers,298        pin_memory=torch.cuda.is_available(),299    )300 301    return train_loader, val_loader, processor302 303 304def build_dataloaders(305    data_dir: Path,306    batch_size: int,307    num_workers: int,308) -> Tuple[DataLoader, DataLoader, object]:309    """Build train and validation dataloaders.310 311    Args:312        data_dir: Path to dataset root.313        batch_size: Batch size.314        num_workers: Number of workers.315 316    Returns:317        Tuple of (train_loader, val_loader, processor).318    """319    # Load processor320    processor = AutoImageProcessor.from_pretrained(MODEL_ID)321 322    # Layout A: split folders (existing format)323    train_dir = data_dir / "train"324    val_dir = data_dir / "valid"325 326    if train_dir.exists() and val_dir.exists():327        train_dataset = DamageDataset(train_dir, processor, is_train=True)328        val_dataset = DamageDataset(val_dir, processor, is_train=False)329 330        train_loader = DataLoader(331            train_dataset,332            batch_size=batch_size,333            shuffle=True,334            num_workers=num_workers,335            pin_memory=torch.cuda.is_available(),336        )337 338        val_loader = DataLoader(339            val_dataset,340            batch_size=batch_size,341            shuffle=False,342            num_workers=num_workers,343            pin_memory=torch.cuda.is_available(),344        )345 346        return train_loader, val_loader, processor347 348    # Layout B: CarDD/FiftyOne damage_assessment format349    assessment_root = _resolve_damage_assessment_root(data_dir)350    if assessment_root is not None:351        return _build_from_damage_assessment(352            data_dir=data_dir,353            batch_size=batch_size,354            num_workers=num_workers,355            processor=processor,356        )357 358    raise ValueError(359        f"Unsupported dataset layout in {data_dir}. Expected either train/valid folders "360        "or a damage_assessment dataset with data/ and samples.json."361    )362