Mini-Nav/mini-nav/simulator/topdown.py

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)


# Default color palette for object labels (cycled if more labels than colors).
_DEFAULT_LABEL_COLORS: dict[str, str] = {
    "chair": "#e6194b",
    "table": "#3cb44b",
    "sofa": "#ffe119",
    "cabinet": "#4363d8",
    "shelf": "#f58231",
    "lamp": "#911eb4",
    "picture": "#42d4f4",
    "window": "#f032e6",
    "door": "#bfef45",
    "plant": "#fabed4",
}

_FALLBACK_COLOR = "#808080"


@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"
    # Object node style.
    object_marker_size: int = 25
    object_label_colors: dict[str, str] = field(
        default_factory=lambda: dict(_DEFAULT_LABEL_COLORS)
    )
    object_fallback_color: str = _FALLBACK_COLOR
    object_show_label: bool = True
    object_label_fontsize: int = 6
    # Edge style.
    edge_color: str = "cyan"
    edge_alpha: float = 0.4
    edge_arrow_size: float = 8.0


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,
    object markers (colored by label), and edges (arrows from room to
    object) 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/object nodes and edges.
        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.
    """
    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 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)

    # --- Room nodes ---
    room_grid_positions: dict[str, tuple[int, int]] = {}
    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,
        )
        room_grid_positions[room_node.room_id] = (grid_x, grid_y)
        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,
        )

    # --- Object nodes (colored by label) ---
    obj_grid_positions: dict[str, tuple[int, int]] = {}
    drawn_labels: set[str] = set()

    for obj_node in elements.object_nodes:
        if obj_node.position is None:
            continue
        grid_y, grid_x = maps.to_grid(
            float(obj_node.position[2]),
            float(obj_node.position[0]),
            top_down_map.shape,
            pathfinder=pathfinder,
        )
        obj_grid_positions[obj_node.obj_id] = (grid_x, grid_y)

        label = obj_node.label or "unknown"
        color = style.object_label_colors.get(label, style.object_fallback_color)

        # Only add to legend once per label.
        legend_label = label if label not in drawn_labels else None
        drawn_labels.add(label)

        ax.scatter(
            grid_x,
            grid_y,
            c=color,
            s=style.object_marker_size,
            marker="^",
            label=legend_label,
            zorder=3,
        )

        if style.object_show_label:
            ax.text(
                grid_x + 1,
                grid_y + 1,
                label,
                color=color,
                fontsize=style.object_label_fontsize,
            )

    # --- Edges (arrows from room → object) ---
    for room_id, obj_id in elements.edges:
        room_pos = room_grid_positions.get(room_id)
        obj_pos = obj_grid_positions.get(obj_id)
        if room_pos is None or obj_pos is None:
            continue
        ax.annotate(
            "",
            xy=obj_pos,
            xytext=room_pos,
            arrowprops=dict(
                arrowstyle="->",
                color=style.edge_color,
                alpha=style.edge_alpha,
                lw=0.8,
            ),
        )

    # --- Legend (only if object labels were drawn) ---
    if drawn_labels:
        ax.legend(
            loc="upper right",
            fontsize=7,
            markerscale=0.8,
            framealpha=0.7,
        )

    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)