mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
from matplotlib.backends.backend_agg import FigureCanvasAgg
|
|
from matplotlib.figure import Figure
|
|
from PIL import Image
|
|
|
|
|
|
@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,
|
|
) -> Image.Image:
|
|
"""Render a top-down scene map as a PIL Image.
|
|
|
|
Uses matplotlib Agg backend (non-interactive) to draw room markers and
|
|
labels on the habitat top-down occupancy map, then returns the result
|
|
as a PIL Image.
|
|
|
|
Args:
|
|
pathfinder: Habitat pathfinder instance.
|
|
elements: Scene elements containing room nodes to annotate.
|
|
meters_per_pixel: Resolution for the top-down map.
|
|
style: Visual style configuration.
|
|
|
|
Returns:
|
|
PIL.Image.Image: Rendered top-down map with annotations.
|
|
|
|
Raises:
|
|
ValueError: If room_nodes is empty or meters_per_pixel <= 0.
|
|
NotImplementedError: If object_nodes or edges are provided.
|
|
"""
|
|
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()
|
|
|
|
# Import habitat maps internally — treated as required dependency.
|
|
from habitat.utils.visualizations import maps
|
|
|
|
# Generate raw top-down occupancy map.
|
|
map_height = float(elements.room_nodes[0].center[1])
|
|
top_down_map = maps.get_topdown_map(
|
|
pathfinder,
|
|
height=map_height,
|
|
meters_per_pixel=meters_per_pixel,
|
|
)
|
|
|
|
# Render with matplotlib Agg backend (no GUI required).
|
|
fig = Figure(figsize=style.figure_size)
|
|
_ = FigureCanvasAgg(fig)
|
|
ax = fig.add_subplot(111)
|
|
ax.imshow(top_down_map, cmap=style.map_cmap)
|
|
|
|
# Annotate room positions and labels.
|
|
for room_node in elements.room_nodes:
|
|
grid_y, grid_x = maps.to_grid(
|
|
float(room_node.center[2]),
|
|
float(room_node.center[0]),
|
|
top_down_map.shape,
|
|
pathfinder=pathfinder,
|
|
)
|
|
ax.scatter(grid_x, grid_y, c=style.room_color, s=style.room_marker_size)
|
|
ax.text(
|
|
grid_x + style.room_label_offset,
|
|
grid_y + style.room_label_offset,
|
|
room_node.room_id,
|
|
color=style.room_label_color,
|
|
fontsize=8,
|
|
)
|
|
|
|
ax.set_title(style.title)
|
|
ax.axis("off")
|
|
fig.canvas.draw()
|
|
|
|
# Extract RGB pixels from the Agg canvas buffer.
|
|
rgba = np.asarray(fig.canvas.buffer_rgba(), dtype=np.uint8)
|
|
rgb = rgba[..., :3].copy()
|
|
|
|
return Image.fromarray(rgb)
|