feat: phase 1 - data pipeline (utils, transforms, dataset)

src/data/dataset.py

@@ -0,0 +1,162 @@
+import os
+from pathlib import Path
+from typing import Optional, Tuple, List
+from PIL import Image
+import torch
+from torch.utils.data import Dataset
+
+from src.data.utils import (
+    parse_voc_xml,
+    get_primary_object,
+    load_image_ids,
+    CLASS_TO_IDX,
+    NUM_CLASSES,
+)
+from src.data.transforms import get_train_transforms, get_val_transforms, Compose
+
+
+class VOCMultiTaskDataset(Dataset):
+    """
+    PASCAL VOC 2012 dataset for multi-task learning.
+
+    Each sample returns:
+        image  : Tensor [3, H, W] — normalized
+        label  : int — primary object class index
+        bbox   : List[float] — [x_min, y_min, x_max, y_max] normalized [0, 1]
+        image_id: str — VOC image ID (e.g. '2007_000032')
+
+    Primary object selection: largest non-difficult bounding box per image.
+    Images with no valid objects are skipped (returns None, filtered by collate_fn).
+    """
+
+    def __init__(
+        self,
+        voc_root: str,
+        split: str = "train",
+        transforms: Optional[Compose] = None,
+        max_samples: Optional[int] = None,
+    ):
+        """
+        Args:
+            voc_root   : path to VOCdevkit/VOC2012/
+            split      : 'train', 'val', or 'trainval'
+            transforms : Compose instance (defaults to split-appropriate transforms)
+            max_samples: cap dataset size — useful for local CPU dev runs
+        """
+        self.voc_root = Path(voc_root)
+        self.split = split
+        self.image_dir = self.voc_root / "JPEGImages"
+        self.annotation_dir = self.voc_root / "Annotations"
+
+        self.image_ids = load_image_ids(str(self.voc_root), split)
+        if max_samples is not None:
+            self.image_ids = self.image_ids[:max_samples]
+
+        if transforms is not None:
+            self.transforms = transforms
+        elif split == "train":
+            self.transforms = get_train_transforms()
+        else:
+            self.transforms = get_val_transforms()
+
+        self._skipped = 0
+
+    def __len__(self) -> int:
+        return len(self.image_ids)
+
+    def __getitem__(self, idx: int):
+        image_id = self.image_ids[idx]
+
+        # --- load annotation ---
+        xml_path = self.annotation_dir / f"{image_id}.xml"
+        try:
+            parsed = parse_voc_xml(str(xml_path))
+        except Exception as e:
+            self._skipped += 1
+            return None
+
+        primary = get_primary_object(parsed)
+        if primary is None:
+            self._skipped += 1
+            return None
+
+        label = primary["label_idx"]
+        bbox  = primary["bbox"]
+
+        # --- load image ---
+        img_path = self.image_dir / f"{parsed['filename']}"
+        if not img_path.exists():
+            # some VOC filenames lack extension
+            img_path = self.image_dir / f"{image_id}.jpg"
+        try:
+            image = Image.open(img_path).convert("RGB")
+        except Exception:
+            self._skipped += 1
+            return None
+
+        # --- apply transforms ---
+        if self.transforms is not None:
+            image, bbox = self.transforms(image, bbox)
+
+        return image, label, bbox, image_id
+
+    def get_class_name(self, idx: int) -> str:
+        from src.data.utils import IDX_TO_CLASS
+        return IDX_TO_CLASS.get(idx, "unknown")
+
+    def class_distribution(self) -> dict:
+        """
+        Iterate all annotations and count primary object per image.
+        Useful for EDA. Slow — don't call during training.
+        """
+        from collections import Counter
+        from src.data.utils import IDX_TO_CLASS
+        counter = Counter()
+        for image_id in self.image_ids:
+            xml_path = self.annotation_dir / f"{image_id}.xml"
+            try:
+                parsed = parse_voc_xml(str(xml_path))
+                primary = get_primary_object(parsed)
+                if primary:
+                    counter[primary["name"]] += 1
+            except Exception:
+                continue
+        return dict(counter)
+
+
+def build_dataloaders(
+    voc_root: str,
+    batch_size: int = 32,
+    num_workers: int = 0,
+    max_train_samples: Optional[int] = None,
+    max_val_samples: Optional[int] = None,
+) -> Tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader]:
+    """
+    Build train and val DataLoaders.
+    num_workers=0 is default for Windows (multiprocessing issues with >0).
+    """
+    from src.data.utils import collate_fn
+
+    train_ds = VOCMultiTaskDataset(
+        voc_root, split="train", max_samples=max_train_samples
+    )
+    val_ds = VOCMultiTaskDataset(
+        voc_root, split="val", max_samples=max_val_samples
+    )
+
+    train_loader = torch.utils.data.DataLoader(
+        train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        collate_fn=collate_fn,
+    )
+    val_loader = torch.utils.data.DataLoader(
+        val_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collate_fn,
+    )
+
+    return train_loader, val_ds

src/data/transforms.py

@@ -0,0 +1,162 @@
+import torch
+import torchvision.transforms.functional as TF
+from torchvision import transforms
+from typing import Tuple, List
+import random
+
+
+class Compose:
+    """Apply a sequence of transforms to both image and bbox."""
+    def __init__(self, transforms_list):
+        self.transforms = transforms_list
+
+    def __call__(self, image, bbox):
+        for t in self.transforms:
+            image, bbox = t(image, bbox)
+        return image, bbox
+
+
+class ToTensor:
+    """Convert PIL image to tensor. Bbox is already a list, pass through."""
+    def __call__(self, image, bbox):
+        return TF.to_tensor(image), bbox
+
+
+class Normalize:
+    """Normalize image tensor. Bbox unchanged."""
+    def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image, bbox):
+        return TF.normalize(image, self.mean, self.std), bbox
+
+
+class Resize:
+    """
+    Resize image to (size, size).
+    Bbox is normalized [0,1] so no change needed.
+    """
+    def __init__(self, size: int = 224):
+        self.size = size
+
+    def __call__(self, image, bbox):
+        image = TF.resize(image, [self.size, self.size])
+        return image, bbox
+
+
+class RandomHorizontalFlip:
+    """
+    Flip image horizontally with probability p.
+    Bbox x-coords must be mirrored: x_min' = 1 - x_max, x_max' = 1 - x_min.
+    """
+    def __init__(self, p: float = 0.5):
+        self.p = p
+
+    def __call__(self, image, bbox):
+        if random.random() < self.p:
+            image = TF.hflip(image)
+            x_min, y_min, x_max, y_max = bbox
+            bbox = [1.0 - x_max, y_min, 1.0 - x_min, y_max]
+        return image, bbox
+
+
+class RandomColorJitter:
+    """Color jitter on image only. Bbox unchanged."""
+    def __init__(self, brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1):
+        self.jitter = transforms.ColorJitter(
+            brightness=brightness,
+            contrast=contrast,
+            saturation=saturation,
+            hue=hue,
+        )
+
+    def __call__(self, image, bbox):
+        return self.jitter(image), bbox
+
+
+class RandomCrop:
+    """
+    Crop a random region of the image, keeping at least min_overlap
+    of the primary bbox inside the crop. Falls back to center crop
+    if a valid crop isn't found within max_attempts.
+
+    Bbox is recalculated relative to the cropped region.
+    """
+    def __init__(self, size: int = 224, min_overlap: float = 0.7, max_attempts: int = 10):
+        self.size = size
+        self.min_overlap = min_overlap
+        self.max_attempts = max_attempts
+
+    def __call__(self, image, bbox):
+        w, h = image.size  # PIL: (width, height)
+        x_min, y_min, x_max, y_max = bbox
+
+        # convert normalized bbox to pixel coords
+        bx1, by1 = x_min * w, y_min * h
+        bx2, by2 = x_max * w, y_max * h
+
+        crop_w = min(self.size, w)
+        crop_h = min(self.size, h)
+
+        for _ in range(self.max_attempts):
+            left = random.randint(0, max(0, w - crop_w))
+            top  = random.randint(0, max(0, h - crop_h))
+            right  = left + crop_w
+            bottom = top  + crop_h
+
+            # intersection with bbox
+            ix1 = max(bx1, left)
+            iy1 = max(by1, top)
+            ix2 = min(bx2, right)
+            iy2 = min(by2, bottom)
+
+            if ix2 > ix1 and iy2 > iy1:
+                inter_area = (ix2 - ix1) * (iy2 - iy1)
+                bbox_area  = (bx2 - bx1) * (by2 - by1)
+                if bbox_area > 0 and (inter_area / bbox_area) >= self.min_overlap:
+                    image = TF.crop(image, top, left, crop_h, crop_w)
+                    image = TF.resize(image, [self.size, self.size])
+
+                    # recalculate bbox relative to crop, re-normalize
+                    new_bbox = [
+                        (bx1 - left) / crop_w,
+                        (by1 - top)  / crop_h,
+                        (bx2 - left) / crop_w,
+                        (by2 - top)  / crop_h,
+                    ]
+                    new_bbox = [max(0.0, min(1.0, v)) for v in new_bbox]
+                    return image, new_bbox
+
+        # fallback: center crop
+        left = (w - crop_w) // 2
+        top  = (h - crop_h) // 2
+        image = TF.crop(image, top, left, crop_h, crop_w)
+        image = TF.resize(image, [self.size, self.size])
+        new_bbox = [
+            (bx1 - left) / crop_w,
+            (by1 - top)  / crop_h,
+            (bx2 - left) / crop_w,
+            (by2 - top)  / crop_h,
+        ]
+        new_bbox = [max(0.0, min(1.0, v)) for v in new_bbox]
+        return image, new_bbox
+
+
+def get_train_transforms(size: int = 224) -> Compose:
+    return Compose([
+        Resize(size + 32),           # resize slightly larger first
+        RandomCrop(size),            # then random crop to target size
+        RandomHorizontalFlip(p=0.5),
+        RandomColorJitter(),
+        ToTensor(),
+        Normalize(),
+    ])
+
+
+def get_val_transforms(size: int = 224) -> Compose:
+    return Compose([
+        Resize(size),
+        ToTensor(),
+        Normalize(),
+    ])

src/data/utils.py

@@ -0,0 +1,137 @@
+import os
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional
+import torch
+from torch.utils.data import DataLoader
+
+
+VOC_CLASSES = [
+    "aeroplane", "bicycle", "bird", "boat", "bottle",
+    "bus", "car", "cat", "chair", "cow",
+    "diningtable", "dog", "horse", "motorbike", "person",
+    "pottedplant", "sheep", "sofa", "train", "tvmonitor"
+]
+
+CLASS_TO_IDX: Dict[str, int] = {cls: idx for idx, cls in enumerate(VOC_CLASSES)}
+IDX_TO_CLASS: Dict[int, str] = {idx: cls for cls, idx in CLASS_TO_IDX.items()}
+NUM_CLASSES = len(VOC_CLASSES)
+
+
+def parse_voc_xml(xml_path: str) -> Dict:
+    """
+    Parse a single PASCAL VOC annotation XML file.
+
+    Returns a dict with:
+        - image_path: str
+        - width: int
+        - height: int
+        - objects: List of dicts, each with:
+            - name: str (class label)
+            - label_idx: int
+            - bbox: [x_min, y_min, x_max, y_max] normalized to [0, 1]
+            - difficult: bool
+    """
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    folder = root.findtext("folder", default="VOC2012")
+    filename = root.findtext("filename")
+
+    size = root.find("size")
+    width = int(size.findtext("width"))
+    height = int(size.findtext("height"))
+
+    objects = []
+    for obj in root.findall("object"):
+        name = obj.findtext("name")
+        if name not in CLASS_TO_IDX:
+            continue
+
+        difficult = bool(int(obj.findtext("difficult", default="0")))
+
+        bndbox = obj.find("bndbox")
+        x_min = float(bndbox.findtext("xmin"))
+        y_min = float(bndbox.findtext("ymin"))
+        x_max = float(bndbox.findtext("xmax"))
+        y_max = float(bndbox.findtext("ymax"))
+
+        # normalize to [0, 1]
+        bbox_norm = [
+            x_min / width,
+            y_min / height,
+            x_max / width,
+            y_max / height,
+        ]
+        # clamp to valid range
+        bbox_norm = [max(0.0, min(1.0, v)) for v in bbox_norm]
+
+        objects.append({
+            "name": name,
+            "label_idx": CLASS_TO_IDX[name],
+            "bbox": bbox_norm,
+            "difficult": difficult,
+        })
+
+    return {
+        "filename": filename,
+        "width": width,
+        "height": height,
+        "objects": objects,
+    }
+
+
+def get_primary_object(parsed: Dict) -> Optional[Dict]:
+    """
+    For multi-task training we need one label + one bbox per image.
+    Strategy: pick the largest non-difficult bounding box by area.
+    Falls back to any object if all are marked difficult.
+    """
+    objects = parsed["objects"]
+    if not objects:
+        return None
+
+    non_difficult = [o for o in objects if not o["difficult"]]
+    candidates = non_difficult if non_difficult else objects
+
+    def bbox_area(obj):
+        b = obj["bbox"]
+        return (b[2] - b[0]) * (b[3] - b[1])
+
+    return max(candidates, key=bbox_area)
+
+
+def load_image_ids(voc_root: str, split: str = "train") -> List[str]:
+    """
+    Load image IDs from VOC ImageSets/Main/<split>.txt.
+    split: 'train', 'val', or 'trainval'
+    """
+    split_file = Path(voc_root) / "ImageSets" / "Main" / f"{split}.txt"
+    if not split_file.exists():
+        raise FileNotFoundError(f"Split file not found: {split_file}")
+
+    with open(split_file) as f:
+        ids = [line.strip() for line in f if line.strip()]
+    return ids
+
+
+def collate_fn(batch: List) -> Tuple:
+    """
+    Custom collate for DataLoader.
+    Filters out None samples (images that failed to parse).
+    Returns:
+        images: Tensor [B, C, H, W]
+        labels: Tensor [B] (long)
+        bboxes: Tensor [B, 4] (float, normalized)
+        image_ids: List[str]
+    """
+    batch = [b for b in batch if b is not None]
+    if not batch:
+        return None, None, None, []
+
+    images = torch.stack([b[0] for b in batch])
+    labels = torch.tensor([b[1] for b in batch], dtype=torch.long)
+    bboxes = torch.tensor([b[2] for b in batch], dtype=torch.float32)
+    image_ids = [b[3] for b in batch]
+
+    return images, labels, bboxes, image_ids