Spaces:
Running
Running
| """Habitat navigation and evaluator-private visual inspection.""" | |
| from __future__ import annotations | |
| import math | |
| from pathlib import Path | |
| from typing import Callable | |
| import habitat_sim | |
| import magnum as mn | |
| import numpy as np | |
| from evolvingnav_paper.habitat_utils import make_simulator, path_distance, sensor_specs, set_agent | |
| def visible_fraction_from_masks(actual: np.ndarray, target_only: np.ndarray, semantic_id: int) -> float: | |
| """Visible target pixels divided by its unobstructed projection at the same pose.""" | |
| projected = int(np.count_nonzero(target_only == semantic_id)) | |
| if projected == 0: | |
| return 0.0 | |
| visible = int(np.count_nonzero(actual == semantic_id)) | |
| return min(1.0, visible / projected) | |
| class HabitatInspectionBackend: | |
| def __init__( | |
| self, hssd_root: Path, scene_id: str, navmesh: Path, | |
| public_viewpoints: dict[int, dict], *, gpu: int = 0, | |
| detector: Callable[[np.ndarray, np.ndarray, str], bool] | None = None, | |
| ) -> None: | |
| self.hssd_root = hssd_root | |
| self.scene_id = scene_id | |
| self.navmesh = navmesh | |
| self.gpu = gpu | |
| self.public_viewpoints = public_viewpoints | |
| self.detector = detector | |
| self.real = None | |
| self.target_only = None | |
| self.target_only_object_id = None | |
| self.object_id = None | |
| self.truth = None | |
| self.projected_pixels = {} | |
| self.last_observation = None | |
| self.last_detections = () | |
| self.last_semantic = None | |
| def distance(self, start, goal) -> float: | |
| if self.real is None: | |
| raise RuntimeError("prepare an episode before requesting paths") | |
| return path_distance(self.real.pathfinder, start, goal) | |
| def _place_target(self, sim, truth: dict, target: dict) -> int: | |
| config_path = self.hssd_root / target["template_handle"] | |
| manager = sim.get_object_template_manager() | |
| manager.load_configs(str(config_path.resolve())) | |
| handles = manager.get_template_handles(target["template_hash"]) | |
| if not handles: | |
| raise RuntimeError(f"target template unavailable: {config_path}") | |
| obj = sim.get_rigid_object_manager().add_object_by_template_handle(handles[-1]) | |
| obj.motion_type = habitat_sim.physics.MotionType.KINEMATIC | |
| obj.semantic_id = self.semantic_id | |
| obj.translation = mn.Vector3(*truth["target_position_xyz"]) | |
| obj.rotation = mn.Quaternion.rotation( | |
| mn.Rad((self.semantic_id % 12) * math.pi / 6), mn.Vector3.y_axis() | |
| ) | |
| sim.perform_discrete_collision_detection() | |
| return obj.object_id | |
| def prepare(self, truth: dict, target: dict, candidate_ids: list[int], | |
| *, dynamic: bool = False) -> None: | |
| self.clear() | |
| if self.real is not None: | |
| self.real.close() | |
| self.real = None | |
| self.truth = dict(truth) | |
| self.semantic_id = int(target["semantic_instance_id"]) | |
| self.target_category = str(target["category_canonical"]) | |
| config = habitat_sim.SimulatorConfiguration() | |
| config.scene_id = "NONE" | |
| config.create_renderer = True | |
| config.enable_physics = True | |
| config.gpu_device_id = self.gpu | |
| agent = habitat_sim.agent.AgentConfiguration() | |
| agent.sensor_specifications = sensor_specs(320, 240) | |
| target_only = habitat_sim.Simulator(habitat_sim.Configuration(config, [agent])) | |
| try: | |
| target_only_id = self._place_target(target_only, truth, target) | |
| self.projected_pixels = {} | |
| for state in candidate_ids: | |
| viewpoint = self.public_viewpoints[state] | |
| set_agent( | |
| target_only.get_agent(0), viewpoint["position_xyz"], | |
| viewpoint["rotation_xyzw"], | |
| ) | |
| semantic = np.asarray(target_only.get_sensor_observations()["semantic"]) | |
| self.projected_pixels[state] = int(np.count_nonzero(semantic == self.semantic_id)) | |
| finally: | |
| if not dynamic: | |
| target_only.close() | |
| if dynamic: | |
| self.target_only = target_only | |
| self.target_only_object_id = target_only_id | |
| self.real = make_simulator( | |
| self.hssd_root, self.scene_id, self.navmesh, self.gpu, 320, 240 | |
| ) | |
| self.object_id = self._place_target(self.real, truth, target) | |
| def inspect(self, state: int, position) -> dict: | |
| if self.truth is None: | |
| raise RuntimeError("prepare an episode before inspection") | |
| viewpoint = self.public_viewpoints[state] | |
| set_agent(self.real.get_agent(0), position, viewpoint["rotation_xyzw"]) | |
| observations = self.real.get_sensor_observations() | |
| self.last_observation = { | |
| "rgb": np.asarray(observations["rgb"]), | |
| "depth": np.asarray(observations["depth"]), | |
| } | |
| actual = np.asarray(observations["semantic"]) | |
| self.last_semantic = actual | |
| actual_pixels = int(np.count_nonzero(actual == self.semantic_id)) | |
| if self.target_only is not None: | |
| set_agent(self.target_only.get_agent(0), position, viewpoint["rotation_xyzw"]) | |
| projection = np.asarray(self.target_only.get_sensor_observations()["semantic"]) | |
| projected_pixels = int(np.count_nonzero(projection == self.semantic_id)) | |
| else: | |
| projected_pixels = self.projected_pixels[state] | |
| visible_fraction = min(1.0, actual_pixels / projected_pixels) if projected_pixels else 0.0 | |
| goal_distance = min( | |
| self.distance(position, goal["position_xyz"]) | |
| for goal in self.truth["valid_goal_viewpoints"] | |
| ) | |
| if self.detector is not None and hasattr(self.detector, "detect_instances"): | |
| self.last_detections = self.detector.detect_instances( | |
| np.asarray(observations["rgb"]), self.target_category | |
| ) | |
| detected = bool(self.last_detections) | |
| elif self.detector is not None: | |
| self.last_detections = () | |
| detected = self.detector( | |
| np.asarray(observations["rgb"]), | |
| np.asarray(observations["depth"]), | |
| self.target_category, | |
| ) | |
| else: | |
| self.last_detections = () | |
| detected = actual_pixels > 0 | |
| return { | |
| "detected": detected, | |
| "visible_fraction": visible_fraction, | |
| "distance_to_valid_goal_m": goal_distance, | |
| "visible_target_pixels": actual_pixels, | |
| "projected_target_pixels": projected_pixels, | |
| } | |
| def move_target(self, event: dict) -> None: | |
| if self.real is None or self.object_id is None: | |
| raise RuntimeError("prepare an episode before moving its target") | |
| position = mn.Vector3(*event["target_position_xyz"]) | |
| self.real.get_rigid_object_manager().get_object_by_id(self.object_id).translation = position | |
| if self.target_only is not None and self.target_only_object_id is not None: | |
| self.target_only.get_rigid_object_manager().get_object_by_id( | |
| self.target_only_object_id | |
| ).translation = position | |
| self.real.perform_discrete_collision_detection() | |
| self.truth["target_position_xyz"] = event["target_position_xyz"] | |
| if "valid_goal_viewpoints" in event: | |
| self.truth["valid_goal_viewpoints"] = event["valid_goal_viewpoints"] | |
| if "current_state_id" in event: | |
| self.truth["current_state_id"] = event["current_state_id"] | |
| def clear(self) -> None: | |
| if self.real is not None and self.object_id is not None: | |
| self.real.get_rigid_object_manager().remove_object_by_id(self.object_id) | |
| self.object_id = None | |
| self.truth = None | |
| self.last_observation = None | |
| self.last_detections = () | |
| self.last_semantic = None | |
| if self.target_only is not None: | |
| self.target_only.close() | |
| self.target_only = None | |
| self.target_only_object_id = None | |
| def close(self) -> None: | |
| self.clear() | |
| if self.real is not None: | |
| self.real.close() | |
| self.real = None | |