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

from __future__ import annotations

import copy
from dataclasses import dataclass, field
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]]
ProgressTrack = Callable[[Iterable[Any], str], Iterable[Any]]


@dataclass(frozen=True)
class RoomView:
    room_id: str
    view_idx: int
    rgb: Any = field(compare=False)
    depth: Any | None = field(default=None, compare=False)
    agent_position: np.ndarray | None = field(default=None, compare=False)
    agent_rotation: Any | None = field(default=None, compare=False)
    camera_position: Any | None = field(default=None, compare=False)
    camera_rotation: Any | None = field(default=None, compare=False)


def flatten_room_views(room_views_by_room: RoomViewsByRoom) -> list[RoomView]:
    result: list[RoomView] = []
    for room_id, views in room_views_by_room.items():
        for idx, item in enumerate(views):
            if isinstance(item, RoomView):
                result.append(item)
            else:
                result.append(RoomView(room_id=room_id, view_idx=idx, rgb=item))
    return result


def _snapshot_rotation(rotation: Any) -> Any:
    """Create an independent copy of *rotation* to prevent mutation aliasing.

    If the value exposes a ``.copy()`` method (e.g. numpy arrays) it is
    preferred; otherwise a stdlib shallow copy is used as a fallback.
    """
    if hasattr(rotation, "copy") and callable(rotation.copy):
        return rotation.copy()
    return copy.copy(rotation)


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",
    depth_sensor_uuid: str | None = None,
    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 view_idx in range(views_per_room):
            observations = sim.get_sensor_observations()
            current_state = agent.get_state()
            agent_position = np.asarray(current_state.position, dtype=np.float32).copy()

            # Capture camera pose from sensor_states when depth sensor is available
            camera_position = None
            camera_rotation = None
            if depth_sensor_uuid is not None:
                sensor_states = getattr(current_state, "sensor_states", None)
                if sensor_states is not None and depth_sensor_uuid in sensor_states:
                    sensor_state = sensor_states[depth_sensor_uuid]
                    camera_position = np.asarray(sensor_state.position, dtype=np.float32).copy()
                    camera_rotation = _snapshot_rotation(sensor_state.rotation)

            room_view = RoomView(
                room_id=room_node.room_id,
                view_idx=view_idx,
                rgb=observations[sensor_uuid],
                depth=observations[depth_sensor_uuid] if depth_sensor_uuid is not None else None,
                agent_position=agent_position,
                agent_rotation=_snapshot_rotation(current_state.rotation),
                camera_position=camera_position,
                camera_rotation=camera_rotation,
            )
            room_views.append(room_view)
            sim.step(turn_action)

        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