feat(dataset): add synthetic dataset generation and configuration

mini-nav/datasets/loader.py

@@ -0,0 +1,105 @@
+"""Data loaders for synthetic and validation datasets."""
+
+from collections.abc import Iterator
+from pathlib import Path
+
+from PIL import Image
+
+
+class SynthDataset:
+    """Dataset loader for synthesized training images."""
+
+    def __init__(self, synth_dir: Path, annotations_suffix: str = ".txt"):
+        """
+        Initialize the synthetic dataset loader.
+
+        Args:
+            synth_dir: Directory containing synthesized images and annotations
+            annotations_suffix: Suffix for annotation files
+        """
+        self.synth_dir = Path(synth_dir)
+        self.annotations_suffix = annotations_suffix
+
+        # Find all images
+        self.image_files = sorted(self.synth_dir.glob("synth_*.jpg"))
+
+    def __len__(self) -> int:
+        return len(self.image_files)
+
+    def __getitem__(self, idx: int) -> tuple[Image.Image, list[tuple[str, int, int, int, int]]]:
+        """Get a single item.
+
+        Args:
+            idx: Index of the item
+
+        Returns:
+            Tuple of (image, annotations) where annotations is a list of
+            (category, xmin, ymin, xmax, ymax)
+        """
+        img_path = self.image_files[idx]
+        image = Image.open(img_path).convert("RGB")
+
+        # Load annotations
+        anno_path = img_path.with_suffix(self.annotations_suffix)
+        annotations: list[tuple[str, int, int, int, int]] = []
+
+        if anno_path.exists():
+            with open(anno_path, "r") as f:
+                for line in f:
+                    line = line.strip()
+                    if line:
+                        parts = line.split()
+                        if len(parts) == 5:
+                            category = parts[0]
+                            xmin, ymin, xmax, ymax = map(int, parts[1:])
+                            annotations.append((category, xmin, ymin, xmax, ymax))
+
+        return image, annotations
+
+    def __iter__(self) -> Iterator[tuple[Image.Image, list[tuple[str, int, int, int, int]]]]:
+        """Iterate over the dataset."""
+        for i in range(len(self)):
+            yield self[i]
+
+
+class ValDataset:
+    """Dataset loader for validation scene images."""
+
+    def __init__(self, scenes_dir: Path, split: str = "easy"):
+        """
+        Initialize the validation dataset loader.
+
+        Args:
+            scenes_dir: Directory containing scene subdirectories
+            split: Scene split to load ('easy' or 'hard')
+        """
+        self.scenes_dir = Path(scenes_dir)
+        self.split = split
+
+        self.split_dir = self.scenes_dir / split
+        if not self.split_dir.exists():
+            raise ValueError(f"Scene split directory not found: {self.split_dir}")
+
+        # Find all RGB images
+        self.image_files = sorted(self.split_dir.glob("*/rgb_*.jpg"))
+
+    def __len__(self) -> int:
+        return len(self.image_files)
+
+    def __getitem__(self, idx: int) -> tuple[Image.Image, Path]:
+        """Get a single item.
+
+        Args:
+            idx: Index of the item
+
+        Returns:
+            Tuple of (image, scene_path)
+        """
+        img_path = self.image_files[idx]
+        image = Image.open(img_path).convert("RGB")
+        return image, img_path.parent
+
+    def __iter__(self) -> Iterator[tuple[Image.Image, Path]]:
+        """Iterate over the dataset."""
+        for i in range(len(self)):
+            yield self[i]