mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
132 lines
4.2 KiB
Python
132 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
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]]
|
|
|
|
|
|
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
|
|
|
|
|
|
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
|