feat(test): add collect test images notebook and replace BitImageProcessorFast

mini-nav/simulator/__init__.py

@@ -5,7 +5,7 @@ from .habitat import (
 )
 from .image_save import save_object_image, save_room_view
 from .topdown import TopDownRenderStyle, TopDownSceneElements, render_topdown_scene_map
-from .views import RoomViewsByRoom, collect_room_views_by_room
+from .views import RoomViewsByRoom, collect_room_views_by_room, collect_scene_images
 
 __all__ = [
     "HabitatSimulatorConfig",
@@ -14,6 +14,7 @@ __all__ = [
     "TopDownSceneElements",
     "close_habitat_simulator",
     "collect_room_views_by_room",
+    "collect_scene_images",
     "create_habitat_simulator",
     "render_topdown_scene_map",
     "save_object_image",

mini-nav/simulator/views.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 from importlib import import_module
+from pathlib import Path
 from typing import Any, Callable, Iterable, Sequence
 
+import numpy as np
 from rich.progress import track
 
 RoomViewsByRoom = dict[str, list[Any]]
@@ -42,3 +44,88 @@ def collect_room_views_by_room(
         all_room_views[room_node.room_id] = room_views
 
     return all_room_views
+
+
+def collect_scene_images(
+    scene_name: str,
+    scene_path: str,
+    output_dir: Path,
+    *,
+    image_size: int = 1024,
+    views_per_point: int = 12,
+    points_per_scene: int = 5,
+    seed: int = 42,
+) -> int:
+    """Collect RGB images from random navigable points in a scene.
+
+    Creates a Habitat simulator for the given scene, samples random
+    navigable points, and captures rotated views at each point. Images
+    are saved as PNG files under ``output_dir / scene_name / {point:03d} /``.
+
+    Args:
+        scene_name: Identifier used as subdirectory name.
+        scene_path: Path to the Habitat scene dataset file (.glb).
+        output_dir: Root output directory for saved images.
+        image_size: Resolution (width and height) of captured images.
+        views_per_point: Number of views captured at each point.
+        points_per_scene: Number of random points to sample.
+        seed: Seed for pathfinder reproducibility.
+
+    Returns:
+        Number of images successfully collected.
+    """
+    from utils.image import numpy_to_pil
+
+    from .habitat import HabitatSimulatorConfig, create_habitat_simulator
+
+    config = HabitatSimulatorConfig(
+        scene_path=scene_path,
+        image_size=image_size,
+        views_per_room=views_per_point,
+    )
+    sim, agent = create_habitat_simulator(config)
+    sim.pathfinder.seed(seed)
+
+    collected_count = 0
+    try:
+        for point_idx in range(points_per_scene):
+            point = None
+            for _ in range(10):
+                candidate = sim.pathfinder.get_random_navigable_point()
+                candidate = np.asarray(candidate, dtype=np.float32)
+                if not np.isfinite(candidate).all():
+                    continue
+                if not sim.pathfinder.is_navigable(candidate):
+                    continue
+                point = candidate
+                break
+
+            if point is None:
+                print(
+                    f"[WARN] Skip {scene_name} point {point_idx:03d}: "
+                    "no valid navigable point"
+                )
+                continue
+
+            agent_state = agent.get_state()
+            agent_state.position = point
+            agent.set_state(agent_state)
+
+            for view_idx in range(views_per_point):
+                obs = sim.get_sensor_observations()
+                rgb = obs["color_sensor"]
+                image = numpy_to_pil(rgb)
+                save_path = (
+                    output_dir
+                    / scene_name
+                    / f"{point_idx:03d}"
+                    / f"view_{view_idx:03d}.png"
+                )
+                save_path.parent.mkdir(parents=True, exist_ok=True)
+                image.save(str(save_path))
+                collected_count += 1
+                sim.step("turn_left")
+    finally:
+        sim.close()
+
+    return collected_count