refactor(simulator): extract habitat simulator and visualization modules

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