mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
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
|