refactor(simulator): extract habitat simulator and visualization modules

mini-nav/simulator/__init__.py

@@ -0,0 +1,18 @@
+from .habitat import (
+    HabitatSimulatorConfig,
+    close_habitat_simulator,
+    create_habitat_simulator,
+)
+from .topdown import TopDownRenderStyle, TopDownSceneElements, render_topdown_scene_map
+from .views import RoomViewsByRoom, collect_room_views_by_room
+
+__all__ = [
+    "HabitatSimulatorConfig",
+    "TopDownRenderStyle",
+    "TopDownSceneElements",
+    "RoomViewsByRoom",
+    "close_habitat_simulator",
+    "collect_room_views_by_room",
+    "create_habitat_simulator",
+    "render_topdown_scene_map",
+]

mini-nav/simulator/habitat.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib import import_module
+from typing import Any
+
+
+@dataclass(frozen=True)
+class HabitatSimulatorConfig:
+    scene_path: str
+    views_per_room: int = 6
+    image_size: int = 256
+    sensor_height: float = 1.5
+    move_forward_step: float = 0.25
+    enable_physics: bool = False
+    sensor_uuid: str = "color_sensor"
+    agent_id: int = 0
+
+
+def create_habitat_simulator(
+    config: HabitatSimulatorConfig,
+    habitat_sim_module: Any | None = None,
+) -> tuple[Any, Any]:
+    if config.views_per_room <= 0:
+        raise ValueError("views_per_room must be greater than 0")
+
+    if config.image_size <= 0:
+        raise ValueError("image_size must be greater than 0")
+
+    if config.move_forward_step <= 0:
+        raise ValueError("move_forward_step must be greater than 0")
+
+    if habitat_sim_module is None:
+        habitat_sim_module = import_module("habitat_sim")
+
+    sim_cfg = habitat_sim_module.SimulatorConfiguration()
+    sim_cfg.scene_id = config.scene_path
+    sim_cfg.enable_physics = config.enable_physics
+
+    agent_cfg = habitat_sim_module.agent.AgentConfiguration()
+    rgb_sensor_spec = habitat_sim_module.CameraSensorSpec()
+    rgb_sensor_spec.uuid = config.sensor_uuid
+    rgb_sensor_spec.sensor_type = habitat_sim_module.SensorType.COLOR
+    rgb_sensor_spec.resolution = [config.image_size, config.image_size]
+    rgb_sensor_spec.position = [0.0, config.sensor_height, 0.0]
+    agent_cfg.sensor_specifications = [rgb_sensor_spec]
+
+    turn_angle = 360.0 / config.views_per_room
+    agent_cfg.action_space = {
+        "move_forward": habitat_sim_module.agent.ActionSpec(
+            "move_forward",
+            habitat_sim_module.agent.ActuationSpec(amount=config.move_forward_step),
+        ),
+        "turn_left": habitat_sim_module.agent.ActionSpec(
+            "turn_left",
+            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
+        ),
+        "turn_right": habitat_sim_module.agent.ActionSpec(
+            "turn_right",
+            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
+        ),
+    }
+
+    simulator_cfg = habitat_sim_module.Configuration(sim_cfg, [agent_cfg])
+    simulator = habitat_sim_module.Simulator(simulator_cfg)
+    agent = simulator.initialize_agent(config.agent_id)
+    return simulator, agent
+
+
+def close_habitat_simulator(simulator: Any) -> None:
+    close = getattr(simulator, "close", None)
+    if callable(close):
+        close()

mini-nav/simulator/topdown.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from importlib import import_module
+from typing import Any, Sequence
+
+
+@dataclass(frozen=True)
+class TopDownSceneElements:
+    room_nodes: Sequence[Any]
+    object_nodes: Sequence[Any] = field(default_factory=tuple)
+    edges: Sequence[tuple[str, str]] = field(default_factory=tuple)
+
+
+@dataclass(frozen=True)
+class TopDownRenderStyle:
+    room_color: str = "red"
+    room_label_color: str = "yellow"
+    room_marker_size: int = 50
+    room_label_offset: int = 2
+    figure_size: tuple[int, int] = (8, 8)
+    map_cmap: str = "gray"
+    title: str = "RoomNode Top-Down Map"
+
+
+def render_topdown_scene_map(
+    pathfinder: Any,
+    elements: TopDownSceneElements,
+    meters_per_pixel: float,
+    style: TopDownRenderStyle | None = None,
+    maps_module: Any | None = None,
+    plt_module: Any | None = None,
+) -> Any:
+    if not elements.room_nodes:
+        raise ValueError("room_nodes must not be empty")
+
+    if meters_per_pixel <= 0:
+        raise ValueError("meters_per_pixel must be greater than 0")
+
+    if elements.object_nodes:
+        raise NotImplementedError("object_nodes overlay is not implemented yet")
+
+    if elements.edges:
+        raise NotImplementedError("edge overlay is not implemented yet")
+
+    if style is None:
+        style = TopDownRenderStyle()
+
+    if maps_module is None:
+        maps_module = import_module("habitat.utils.visualizations.maps")
+
+    if plt_module is None:
+        plt_module = import_module("matplotlib.pyplot")
+
+    map_height = float(elements.room_nodes[0].center[1])
+    top_down_map = maps_module.get_topdown_map(
+        pathfinder,
+        height=map_height,
+        meters_per_pixel=meters_per_pixel,
+    )
+
+    plt_module.figure(figsize=style.figure_size)
+    plt_module.imshow(top_down_map, cmap=style.map_cmap)
+
+    for room_node in elements.room_nodes:
+        grid_y, grid_x = maps_module.to_grid(
+            float(room_node.center[2]),
+            float(room_node.center[0]),
+            top_down_map.shape,
+            pathfinder=pathfinder,
+        )
+        plt_module.scatter(grid_x, grid_y, c=style.room_color, s=style.room_marker_size)
+        plt_module.text(
+            grid_x + style.room_label_offset,
+            grid_y + style.room_label_offset,
+            room_node.room_id,
+            color=style.room_label_color,
+            fontsize=8,
+        )
+
+    plt_module.title(style.title)
+    plt_module.axis("off")
+    plt_module.show()
+    return top_down_map

mini-nav/simulator/views.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any, Callable, Iterable, Sequence
+
+from rich.progress import track
+
+RoomViewsByRoom = dict[str, list[Any]]
+ProgressTrack = Callable[[Iterable[Any], str], Iterable[Any]]
+
+
+def collect_room_views_by_room(
+    agent: Any,
+    sim: Any,
+    room_nodes: Sequence[Any],
+    views_per_room: int,
+    *,
+    habitat_sim_module: Any | None = None,
+    sensor_uuid: str = "color_sensor",
+    turn_action: str = "turn_left",
+    progress_description: str = "Collecting room views",
+    progress_track: ProgressTrack = track,
+) -> RoomViewsByRoom:
+    if views_per_room <= 0:
+        raise ValueError("views_per_room must be greater than 0")
+
+    if habitat_sim_module is None:
+        habitat_sim_module = import_module("habitat_sim")
+
+    all_room_views: RoomViewsByRoom = {}
+    for room_node in progress_track(room_nodes, progress_description):
+        agent_state = habitat_sim_module.AgentState()
+        agent_state.position = room_node.center.copy()
+        agent.set_state(agent_state)
+
+        room_views = []
+        for _ in range(views_per_room):
+            observations = sim.get_sensor_observations()
+            room_views.append(observations[sensor_uuid])
+            sim.step(turn_action)
+
+        all_room_views[room_node.room_id] = room_views
+
+    return all_room_views