Add HF-compatible dataset loader

dataset.py

@@ -0,0 +1,203 @@
+"""
+IPAD Dataset Loader for HuggingFace Infrastructure
+Loads data from HF Hub and provides PyTorch DataLoader compatible interface
+"""
+import torch
+from torch.utils.data import Dataset, DataLoader
+import cv2
+import numpy as np
+from pathlib import Path
+import zipfile
+from huggingface_hub import hf_hub_download
+import os
+from typing import List, Tuple, Optional
+import random
+
+class IPADVideoDataset(Dataset):
+    """
+    IPAD Video Anomaly Detection Dataset
+
+    Args:
+        root_dir: Path to extracted dataset
+        device_name: Device ID (e.g., "S01", "S02", ..., "S12")
+        split: "train" or "test"
+        clip_length: Number of frames per clip (default: 16)
+        frame_size: Tuple of (height, width) for resizing (default: (256, 256))
+        stride: Frame sampling stride (default: 1)
+        normalize: Whether to normalize frames to [-1, 1]
+    """
+
+    def __init__(
+        self,
+        root_dir: str,
+        device_name: str = "S01",
+        split: str = "train",
+        clip_length: int = 16,
+        frame_size: Tuple[int, int] = (256, 256),
+        stride: int = 1,
+        normalize: bool = True
+    ):
+        self.root_dir = Path(root_dir)
+        self.device_name = device_name
+        self.split = split
+        self.clip_length = clip_length
+        self.frame_size = frame_size
+        self.stride = stride
+        self.normalize = normalize
+
+        # Construct path to device frames
+        self.device_path = self.root_dir / device_name / split / "frames"
+
+        if not self.device_path.exists():
+            raise ValueError(f"Dataset path not found: {self.device_path}")
+
+        # Get all video directories
+        self.video_dirs = sorted([d for d in self.device_path.iterdir() if d.is_dir()])
+
+        # Build index of all valid clips
+        self.clips = []
+        for video_dir in self.video_dirs:
+            frames = sorted(list(video_dir.glob("*.jpg")) + list(video_dir.glob("*.png")))
+            num_frames = len(frames)
+
+            # Create clips with stride
+            for start_idx in range(0, num_frames - clip_length + 1, stride):
+                self.clips.append({
+                    'video_dir': video_dir,
+                    'start_idx': start_idx,
+                    'frames': frames[start_idx:start_idx + clip_length]
+                })
+
+        print(f"Loaded {len(self.clips)} clips from {device_name}/{split}")
+
+    def __len__(self) -> int:
+        return len(self.clips)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        clip_info = self.clips[idx]
+        frames = []
+
+        # Load and process each frame
+        for frame_path in clip_info['frames']:
+            frame = cv2.imread(str(frame_path))
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = cv2.resize(frame, self.frame_size)
+
+            # Normalize to [0, 1]
+            frame = frame.astype(np.float32) / 255.0
+
+            # Normalize to [-1, 1] if requested
+            if self.normalize:
+                frame = (frame - 0.5) / 0.5
+
+            frames.append(frame)
+
+        # Convert to tensor: [T, H, W, C] -> [C, T, H, W]
+        frames = np.stack(frames, axis=0)  # [T, H, W, C]
+        frames = torch.from_numpy(frames).permute(3, 0, 1, 2)  # [C, T, H, W]
+
+        return frames
+
+
+def download_and_extract_dataset(cache_dir: str = "./cache") -> Path:
+    """
+    Download IPAD dataset from HF Hub and extract it
+
+    Returns:
+        Path to extracted dataset directory
+    """
+    cache_dir = Path(cache_dir)
+    cache_dir.mkdir(exist_ok=True, parents=True)
+
+    extracted_path = cache_dir / "ipad_dataset"
+
+    if extracted_path.exists():
+        print(f"✅ Dataset already extracted at {extracted_path}")
+        return extracted_path
+
+    print("📥 Downloading dataset from HF Hub...")
+    zip_path = hf_hub_download(
+        repo_id="MSherbinii/ipad-industrial-anomaly",
+        filename="ipad_dataset.zip",
+        repo_type="dataset",
+        cache_dir=str(cache_dir)
+    )
+
+    print("📦 Extracting dataset...")
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(cache_dir)
+
+    print(f"✅ Dataset extracted to {extracted_path}")
+    return extracted_path
+
+
+def create_dataloaders(
+    dataset_path: str,
+    device_name: str = "S01",
+    batch_size: int = 4,
+    num_workers: int = 4,
+    clip_length: int = 16,
+    frame_size: Tuple[int, int] = (256, 256)
+) -> Tuple[DataLoader, DataLoader]:
+    """
+    Create train and test DataLoaders for a specific device
+
+    Args:
+        dataset_path: Path to extracted IPAD dataset
+        device_name: Device ID (e.g., "S01")
+        batch_size: Batch size for DataLoader
+        num_workers: Number of worker processes
+        clip_length: Frames per clip
+        frame_size: Frame dimensions
+
+    Returns:
+        Tuple of (train_loader, test_loader)
+    """
+    train_dataset = IPADVideoDataset(
+        root_dir=dataset_path,
+        device_name=device_name,
+        split="train",
+        clip_length=clip_length,
+        frame_size=frame_size,
+        stride=clip_length // 2  # 50% overlap for training
+    )
+
+    test_dataset = IPADVideoDataset(
+        root_dir=dataset_path,
+        device_name=device_name,
+        split="test",
+        clip_length=clip_length,
+        frame_size=frame_size,
+        stride=clip_length  # No overlap for testing
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=True
+    )
+
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=False
+    )
+
+    return train_loader, test_loader
+
+
+# Device name mappings
+DEVICE_NAMES = [
+    "S01", "S02", "S03", "S04", "S05", "S06",
+    "S07", "S08", "S09", "S10", "S11", "S12",
+    "R01", "R02", "R03", "R04"
+]
+
+SYNTHETIC_DEVICES = [f"S{i:02d}" for i in range(1, 13)]
+REAL_DEVICES = [f"R{i:02d}" for i in range(1, 5)]