diff --git a/README.md b/README.md index e6015f7..3b468fd 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ data-sampling approaches, typical RL terms and a benchmarking environment. Curre This extension lets you load realistic terrains complete with rich semantic annotations, run fast traversability analysis, and render large batches of multi-modal data. It exposes three core modules: - **Environment Importer** – load Matterport, Unreal/Carla, generated or USD terrains and expose all geometric / semantic domains → [Details](exts/nav_suite/docs/README.md#environment-importer) - - **Data Collectors** – sample trajectories, viewpoints and render multi-modal data from any imported world → [Details](exts/nav_suite/docs/README.md#data-collectors) + - **Data Collectors** – sample trajectories and render multi-modal sensor data from any imported world → [Details](exts/nav_suite/docs/README.md#data-collectors) - **Terrain Analysis** – build traversability height-maps and graphs for path planning and curriculum tasks → [Details](exts/nav_suite/docs/README.md#traversabilty-analysis) ## `nav_tasks` Extension @@ -177,7 +177,7 @@ Here we provide a set of examples that demonstrate how to use the different part - [Import the Nvidia Warehouse Environment](scripts/nav_suite/terrains/warehouse_import.py) - ``collector`` - [Sample Trajectories from Matterport](scripts/nav_suite/collector/matterport_trajectory_sampling.py) - - [Sample Viewpoints and Render Images from Carla (Unreal Engine)](scripts/nav_suite/collector/carla_viewpoint_sampling.py) + - [Sample Camera Viewpoints and Render Images from Carla (Unreal Engine)](scripts/nav_suite/collector/carla_sensor_data_sampling.py) ## Citing diff --git a/exts/nav_suite/config/extension.toml b/exts/nav_suite/config/extension.toml index b97c8f5..18610ee 100644 --- a/exts/nav_suite/config/extension.toml +++ b/exts/nav_suite/config/extension.toml @@ -1,7 +1,7 @@ [package] # Note: Semantic Versioning is used: https://semver.org/ -version = "0.2.4" +version = "0.2.5" # Description title = "IsaacLab Navigation Suite" diff --git a/exts/nav_suite/docs/CHANGELOG.rst b/exts/nav_suite/docs/CHANGELOG.rst index 1338065..868c4a5 100644 --- a/exts/nav_suite/docs/CHANGELOG.rst +++ b/exts/nav_suite/docs/CHANGELOG.rst @@ -1,6 +1,27 @@ Changelog --------- + +0.2.5 (2025-08-13) +~~~~~~~~~~~~~~~~~~ + +Added +^^^^^ + +- Added sampling support to other sensor data (such as RayCasters) to :class:`nav_suite.collectors.SensorDataSampling` + for sampling sensor data from the environment + - RayCaster implementation in :class:`nav_suite.collectors.sensors.RayCasterSensor` + +Changed +^^^^^^^ + +- Renamed :class:`nav_suite.collectors.ViewpointSampling` to :class:`nav_suite.collectors.SensorSampling` and + :class:`nav_suite.collectors.ViewpointSamplingCfg` to :class:`nav_suite.collectors.SensorSamplingCfg` +- Sensor data extraction is now done in individual classes for each sensor type. The logic for camera data prev. + included in :class:`nav_suite.collectors.ViewpointSampling` is now extracted in + :class:`nav_suite.collectors.sensors.CameraSensor`. + + 0.2.4 (2025-08-08) ~~~~~~~~~~~~~~~~~~ diff --git a/exts/nav_suite/docs/README.md b/exts/nav_suite/docs/README.md index 7a85cae..c8cdcc6 100644 --- a/exts/nav_suite/docs/README.md +++ b/exts/nav_suite/docs/README.md @@ -92,35 +92,21 @@ This extensions allows to collect data from the previously loaded environments a The trajectory sampling can be executed multiple times with different number of sampled trajectories as well as different minimum and maximum lengths. -- `Viewpoint sampling and image rendering`: +- `Sensor data sampling and rendering`: - For the viewpoint sampling the same terrain analysis as for the trajectory sampling is executed. The graph and traversability parameters are defined in corresponding [config file](../nav_suite/terrain_analysis/terrain_analysis_cfg.py). + For the sensor data sampling the same terrain analysis as for the trajectory sampling is executed. The graph and traversability parameters are defined in corresponding [config file](../nav_suite/terrain_analysis/terrain_analysis_cfg.py). **Important** for the analysis also regarding the semantic domain, a semantic class to cost mapping has to be defined in the config. Per default, an example cost map for ``matterport`` environments is selected. - Each node of the prah is a possible viewpoint, with the orientation uniformly sampled between variable bounds. The exact parameters of the sampling can be defined [here](../nav_suite/collectors/viewpoint_sampling_cfg.py). You can define the ``module`` and the ``class`` of the parameters config that is used for the sampling. An example is provided that is optimized for the legged robot ANYmal and a matterport environment. Please not that this configuration assumes that two cameras are added where the first one has access to semantic information and the second to geoemtric information. - - The number of viepoints that are sampled can be directory defined in the GUI. With the button ``Viewpoint Sampling`` the viewpoints are saved as ``camera_poses`` under the defined directory. Afterwards, click ``Viewpoint Renedering`` to get the final rendered images. The resulting folder structure is as follows: - - ``` graphql - cfg.data_dir - ├── camera_poses.txt # format: x y z qw qx qy qz - ├── cfg.depth_cam_name # required - | ├── intrinsics.txt # K-matrix (3x3) - | ├── distance_to_image_plane # annotator - | | ├── xxxx.png # images saved with 4 digits, e.g. 0000.png - ├── cfg.depth_cam_name # optional - | ├── intrinsics.txt # K-matrix (3x3) - | ├── distance_to_image_plane # annotator - | | ├── xxxx.png # images saved with 4 digits, e.g. 0000.png - ``` + Each node of the graph is a possible data sampling point, with the orientation uniformly sampled between variable bounds. The exact parameters of the sampling can be defined [here](../nav_suite/collectors/sensor_data_sampling_cfg.py). How the individual sensor data is treated is defined in individual sensor modules, i.e., [Camera](../nav_suite/collectors/sensors/camera_cfg.py), [RayCaster](../nav_suite/collectors/sensors/raycaster_cfg.py). + An example is provided that is optimized for the legged robot ANYmal and a matterport environment. Please note that this configuration assumes that two cameras are added where the first one has access to semantic information and the second to geometric information. ### Standalone scripts Standalone scripts are provided to demonstrate the loading of different environments: - [Sample Trajectories from Matterport](../../../scripts/nav_suite/collector/matterport_trajectory_sampling.py) - - [Sample Viewpoints and Render Images from Carla (Unreal Engine)](../../../scripts/nav_suite/collector/carla_viewpoint_sampling.py) + - [Sample Camera Viewpoints and Render Images from Carla (Unreal Engine)](../../../scripts/nav_suite/collector/carla_sensor_data_sampling.py) > [!NOTE] **Matterport Sensors**: \ diff --git a/exts/nav_suite/nav_suite/collectors/__init__.py b/exts/nav_suite/nav_suite/collectors/__init__.py index 2aa9d2d..cefc697 100644 --- a/exts/nav_suite/nav_suite/collectors/__init__.py +++ b/exts/nav_suite/nav_suite/collectors/__init__.py @@ -3,7 +3,8 @@ # # SPDX-License-Identifier: Apache-2.0 +from .sensor_data_sampling import SensorDataSampling +from .sensor_data_sampling_cfg import SensorDataSamplingCfg +from .sensors import * # noqa: F401, F403 from .trajectory_sampling import TrajectorySampling from .trajectory_sampling_cfg import TrajectorySamplingCfg -from .viewpoint_sampling import ViewpointSampling -from .viewpoint_sampling_cfg import ViewpointSamplingCfg diff --git a/exts/nav_suite/nav_suite/collectors/sensor_data_sampling.py b/exts/nav_suite/nav_suite/collectors/sensor_data_sampling.py new file mode 100644 index 0000000..0cded98 --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensor_data_sampling.py @@ -0,0 +1,340 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import os +import pickle +import random +import torch + +import isaaclab.utils.math as math_utils +import omni.log +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import GREEN_ARROW_X_MARKER_CFG +from isaaclab.scene import InteractiveScene +from isaaclab.sim import SimulationContext + +from .sensor_data_sampling_cfg import SensorDataSamplingCfg + + +class SensorDataSampling: + def __init__(self, cfg: SensorDataSamplingCfg, scene: InteractiveScene): + # check the config for any missing values + cfg.validate() + + # save cfg and env + self.cfg = cfg + self.scene = scene + + # get sim context + self.sim = SimulationContext.instance() + + # analyse terrains -- check if singleton is used and available + if ( + hasattr(self.cfg.terrain_analysis.class_type, "instance") + and self.cfg.terrain_analysis.class_type.instance() is not None + ): + self.terrain_analyser = self.cfg.terrain_analysis.class_type.instance() + else: + self.terrain_analyser = self.cfg.terrain_analysis.class_type(self.cfg.terrain_analysis, scene=self.scene) + + # init the sensor data handlers + self.sensor_data_handlers = [ + sensor_cfg.class_type(sensor_cfg, scene=self.scene) for sensor_cfg in self.cfg.sensor_data_handlers + ] + + ### + # Properties + ### + + @property + def samples(self) -> torch.Tensor | list[torch.Tensor]: + """Get the sample points. + + Samples are stored in a torch tensor with the structure + [x, y, z, qw, qx, qv, qz] + """ + return self._samples + + @property + def sliced_bounding_boxes(self) -> list[tuple[float, float, float, float]] | None: + """Get the sliced bounding boxes. + + The sliced bounding boxes follow the terrain analysis format of the mesh dimensions, i.e. [x_max, y_max, x_min, y_min] + """ + if self.cfg.sliced_sampling is not None and hasattr(self, "_sliced_bounding_boxes"): + return self._sliced_bounding_boxes + else: + return None + + ### + # Operations + ### + + def sample_sensor_data(self, nbr_samples: int, seed: int = 1) -> torch.Tensor | list[torch.Tensor]: + """Sample sensor data for the given number of samples and seed. + + Samples are stored in a torch tensor with the structure + [x, y, z, qw, qx, qv, qz] + + Args: + nbr_samples (int): The number of sample points. + seed (int, optional): The seed for the random number generator. Defaults to 1. + + Returns: + torch.Tensor | list[torch.Tensor]: The sampled points. + """ + + # get the number of slices and their locations + # NOTE: the number of sample points is kept constant per slice + if self.cfg.sliced_sampling is not None: + # get mesh size from terrain analysis -- [x_max, y_max, x_min, y_min] + mesh_size = self.terrain_analyser.mesh_dimensions + # get the number of slices in both x and y direction + nbr_slices_x = int(np.ceil((mesh_size[0] - mesh_size[2]) / self.cfg.sliced_sampling[0])) + nbr_slices_y = int(np.ceil((mesh_size[1] - mesh_size[3]) / self.cfg.sliced_sampling[1])) + slice_locations_x = np.linspace(mesh_size[2], mesh_size[0], nbr_slices_x) + slice_locations_y = np.linspace(mesh_size[3], mesh_size[1], nbr_slices_y) + # get the slice bounding boxes + # slide bounding boxes follow the terrain analysis format of the mesh dimensions, i.e. [x_max, y_max, x_min, y_min] + self._sliced_bounding_boxes = [] + for i in range(nbr_slices_x - 1): + for j in range(nbr_slices_y - 1): + self._sliced_bounding_boxes.append( + (slice_locations_x[i + 1], slice_locations_y[j + 1], slice_locations_x[i], slice_locations_y[j]) + ) + + # execute sensor data sampling for each slice + self._samples = [] + failed_slices = [] + for i, slice_bounding_box in enumerate(self._sliced_bounding_boxes): + try: + self._samples.append(self._sample_point_per_area(nbr_samples, seed, slice_bounding_box)) + except Exception as e: + failed_slices.append(slice_bounding_box) + omni.log.warn(f"Error sampling sensor data for slice {slice_bounding_box}: {e}") + continue + + # remove failed slices + self._sliced_bounding_boxes = [box for box in self._sliced_bounding_boxes if box not in failed_slices] + # dump the sliced bounding boxes and samples + filedir = self.cfg.save_path if self.cfg.save_path else self._get_save_filedir() + with open(os.path.join(filedir, "sliced_bounding_boxes.pkl"), "wb") as f: + pickle.dump(self._sliced_bounding_boxes, f) + + else: + # execute sample point sampling for the whole mesh + self._samples = self._sample_point_per_area(nbr_samples, seed) + + return self._samples + + def render_sensor_data( + self, samples: torch.Tensor | list[torch.Tensor] | None = None + ) -> torch.Tensor | list[torch.Tensor]: + """Render the sensor data at the given sample points and save them to the drive.""" + if samples is None: + samples = self.samples + + if isinstance(samples, list): + assert ( + self.sliced_bounding_boxes is not None + ), "Sliced bounding boxes must be set to render a list of sample points." + assert len(samples) == len( + self.sliced_bounding_boxes + ), "The number of samples must match the number of sliced bounding boxes." + render_sample_points = [] + for sample, slice_bounding_box in zip(samples, self.sliced_bounding_boxes): + print(f"[INFO] Rendering sample points for slice {slice_bounding_box}") + render_sample_points.append(self._render_sensor_data_per_area(sample, slice_bounding_box)) + else: + render_sample_points = self._render_sensor_data_per_area(samples) + + return render_sample_points + + ### + # Helper functions + ### + + def _sample_point_per_area( + self, nbr_samples: int, seed: int = 1, slice_bounding_box: tuple[float, float, float, float] | None = None + ) -> torch.Tensor: + + filedir = self.cfg.save_path if self.cfg.save_path else self._get_save_filedir() + if slice_bounding_box is not None: + slice_bounding_box_str = "slice_" + "_".join( + f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in slice_bounding_box + ) + filedir = os.path.join(filedir, slice_bounding_box_str) + elif self.cfg.terrain_analysis.terrain_bounding_box is not None: + slice_bounding_box_str = "slice_" + "_".join( + f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in self.cfg.terrain_analysis.terrain_bounding_box + ) + filedir = os.path.join(filedir, slice_bounding_box_str) + else: + slice_bounding_box_str = "full" + filename = os.path.join(filedir, f"sample_points_seed{seed}_samples{nbr_samples}.pkl") + if os.path.isfile(filename): + with open(filename, "rb") as f: + data = pickle.load(f) + # add loaded path dict to data dict + omni.log.info(f"Loaded {nbr_samples} with seed {seed} for map part {slice_bounding_box_str}.") + return data + else: + omni.log.info( + f"No sample points found for seed {seed} and {nbr_samples} samples for map part" + f" {slice_bounding_box_str}." + ) + + print(f"[INFO] Sampling sample points for {slice_bounding_box_str}") + + # analyse terrain if not done yet + if slice_bounding_box is not None: + # overwrite the mesh dimensions to influence where the sampling is done + self.terrain_analyser.reset_graph() + self.terrain_analyser._mesh_dimensions = slice_bounding_box + self.terrain_analyser.analyse() + elif not self.terrain_analyser.complete: + self.terrain_analyser.analyse() + + # set seed + random.seed(seed) + omni.log.info(f"Start sampling {nbr_samples} sample points for map part {slice_bounding_box_str}.") + + # samples are organized in [point_idx, neighbor_idx, distance] + # sample from each point the neighbor with the largest distance + nbr_samples_per_point = int(np.ceil(nbr_samples / self.terrain_analyser.points.shape[0]).item()) + sample_locations = torch.zeros((nbr_samples_per_point * self.terrain_analyser.points.shape[0], 2)) + sample_locations_count = 0 + curr_point_idx = 0 + while sample_locations_count < nbr_samples: + # get samples + sample_idx = self.terrain_analyser.samples[:, 0] == curr_point_idx + sample_idx_select = torch.randperm(sample_idx.sum())[:nbr_samples_per_point] + sample_locations[sample_locations_count : sample_locations_count + sample_idx_select.shape[0]] = ( + self.terrain_analyser.samples[sample_idx][sample_idx_select, :2] + ) + sample_locations_count += sample_idx_select.shape[0] + curr_point_idx += 1 + # reset point index if all points are sampled + if curr_point_idx >= self.terrain_analyser.points.shape[0]: + curr_point_idx = 0 + + sample_locations = sample_locations[:sample_locations_count].type(torch.int64) + + # get the z angle of the neighbor that is closest to the origin point + neighbor_direction = ( + self.terrain_analyser.points[sample_locations[:, 0]] - self.terrain_analyser.points[sample_locations[:, 1]] + ) + z_angles = torch.atan2(neighbor_direction[:, 1], neighbor_direction[:, 0]).to("cpu") + + # vary the rotation of the forward and horizontal axis (in camera frame) as a uniform distribution within the limits + x_angles = math_utils.sample_uniform( + self.cfg.x_angle_range[0], self.cfg.x_angle_range[1], sample_locations_count, device="cpu" + ) + y_angles = math_utils.sample_uniform( + self.cfg.y_angle_range[0], self.cfg.y_angle_range[1], sample_locations_count, device="cpu" + ) + x_angles = torch.deg2rad(x_angles) + y_angles = torch.deg2rad(y_angles) + + samples = torch.zeros((sample_locations_count, 7)) + samples[:, :3] = self.terrain_analyser.points[sample_locations[:, 0]] + samples[:, 3:] = math_utils.quat_from_euler_xyz(x_angles, y_angles, z_angles) + + omni.log.info(f"Sampled {sample_locations_count} sample points for map part {slice_bounding_box_str}.") + + # save samples + os.makedirs(filedir, exist_ok=True) + with open(filename, "wb") as f: + pickle.dump(samples, f) + + omni.log.info(f"Saved {sample_locations_count} sample points with seed {seed} to {filename}.") + + # debug points and orientation + if self.cfg.debug_viz: + env_render_steps = 1000 + marker_cfg = GREEN_ARROW_X_MARKER_CFG.copy() + marker_cfg.prim_path = "/Visuals/sample_points" + marker_cfg.markers["arrow"].scale = (0.1, 0.1, 0.1) + self.visualizer = VisualizationMarkers(marker_cfg) + self.visualizer.visualize(samples[:, :3], samples[:, 3:]) + + omni.log.info(f"Visualizing {sample_locations_count} samples for {env_render_steps} render steps...") + for _ in range(env_render_steps): + self.sim.render() + + self.visualizer.set_visibility(False) + omni.log.info("Done visualizing.") + + return samples + + def _render_sensor_data_per_area( + self, samples: torch.Tensor, slice_bounding_box: tuple[float, float, float, float] | None = None + ) -> torch.Tensor: + omni.log.info(f"Start rendering sensor data at {samples.shape[0]} points.") + # define how many rounds are necessary to render all sample points + num_rounds = int(np.ceil(samples.shape[0] / self.scene.num_envs)) + + filedir = self.cfg.save_path if self.cfg.save_path else self._get_save_filedir() + if slice_bounding_box is not None: + slice_bounding_box_str = "slice_" + "_".join( + f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in slice_bounding_box + ) + filedir = os.path.join(filedir, slice_bounding_box_str) + elif self.cfg.terrain_analysis.terrain_bounding_box is not None: + slice_bounding_box_str = "slice_" + "_".join( + f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in self.cfg.terrain_analysis.terrain_bounding_box + ) + filedir = os.path.join(filedir, slice_bounding_box_str) + + # pre_collection callback + for sensor_data_handler in self.sensor_data_handlers: + sensor_data_handler.pre_collection(samples, filedir) + + # save images + samples = samples.to(self.scene.device) + for i in range(num_rounds): + # get samples idx + samples_idx = torch.arange(i * self.scene.num_envs, min((i + 1) * self.scene.num_envs, samples.shape[0])) + env_ids = torch.arange(samples_idx.shape[0]) + + # execute pre_sim_update callback + for sensor_data_handler in self.sensor_data_handlers: + sensor_data_handler.pre_sim_update(samples[samples_idx, :3], samples[samples_idx, 3:], env_ids) + + # update simulation + self.scene.write_data_to_sim() + + # perform render steps to fill buffers if needed by any sensor + if any([sensor_data_handler.cfg.requires_render for sensor_data_handler in self.sensor_data_handlers]): + for _ in range(10): + self.sim.render() + + # update scene buffers + self.scene.update(self.sim.get_physics_dt()) + + # execute post_sim_update callback + for sensor_data_handler in self.sensor_data_handlers: + sensor_data_handler.post_sim_update(samples_idx, env_ids, filedir, slice_bounding_box) + + # execute post_collection callback + for sensor_data_handler in self.sensor_data_handlers: + sensor_data_handler.post_collection(samples, filedir) + + return samples + + def _get_save_filedir(self) -> str: + # get env name + if isinstance(self.scene.terrain.cfg.usd_path, str): + terrain_file_path = self.scene.terrain.cfg.usd_path + else: + raise KeyError("Only implemented for terrains loaded from usd and matterport") + env_name = os.path.splitext(terrain_file_path)[0] + # create directory if necessary + filedir = os.path.join(terrain_file_path, env_name) + os.makedirs(filedir, exist_ok=True) + return filedir diff --git a/exts/nav_suite/nav_suite/collectors/viewpoint_sampling_cfg.py b/exts/nav_suite/nav_suite/collectors/sensor_data_sampling_cfg.py similarity index 52% rename from exts/nav_suite/nav_suite/collectors/viewpoint_sampling_cfg.py rename to exts/nav_suite/nav_suite/collectors/sensor_data_sampling_cfg.py index f206021..abe9609 100644 --- a/exts/nav_suite/nav_suite/collectors/viewpoint_sampling_cfg.py +++ b/exts/nav_suite/nav_suite/collectors/sensor_data_sampling_cfg.py @@ -3,68 +3,53 @@ # # SPDX-License-Identifier: Apache-2.0 +from dataclasses import MISSING + from isaaclab.utils import configclass from ..terrain_analysis import TerrainAnalysisCfg, TerrainAnalysisSingletonCfg +from .sensors.base_cfg import SensorBaseCfg @configclass -class ViewpointSamplingCfg: - """Configuration for the viewpoint sampling.""" - - terrain_analysis: TerrainAnalysisCfg | TerrainAnalysisSingletonCfg = TerrainAnalysisCfg(raycaster_sensor="camera_0") - """Name of the camera object in the scene definition used for the terrain analysis.""" +class SensorDataSamplingCfg: + """Configuration for the sensor data sampling.""" ### - # Camera Configuration + # Sensor Configuration ### - cameras: dict[str, list[str]] = { - "camera_0": ["semantic_segmentation"], - "camera_1": ["distance_to_image_plane"], - } - """Dict of cameras and corresponding annotators to use for the viewpoint sampling.""" - depth_scale: float = 1000.0 - """Scaling factor for the depth values.""" + sensor_data_handlers: list[SensorBaseCfg] = MISSING + """List of sensor data handlers that include the logic of how the sensor data is treated.""" ### # Sampling Configuration ### + terrain_analysis: TerrainAnalysisCfg | TerrainAnalysisSingletonCfg = TerrainAnalysisCfg(raycaster_sensor="camera_0") + """Name of the camera object in the scene definition used for the terrain analysis.""" + sample_points: int = 10000 """Number of random points to sample.""" + x_angle_range: tuple[float, float] = (-2.5, 2.5) y_angle_range: tuple[float, float] = (-2, 5) # negative angle means in isaac convention: look down """Range of the x and y angle of the camera (in degrees), will be randomly selected according to a uniform distribution""" + height: float = 0.5 """Height to use for the random points.""" + sliced_sampling: tuple[float, float] | None = None - """Sliced sampling of the viewpoints. If None, no slicing is applied. + """Sliced sampling of the sample points. If None, no slicing is applied. If a tuple is provided, the points will be sampled in slices of the given width and length. That means, the terrain - will be split into slices and for each slice, the number of viewpoints as given to the ``sample_viewpoints`` method - is generated. - """ - - ### - # Point Cloud Generation - ### - generate_point_cloud: bool = True - """Whether to generate a point cloud from the depth images.""" - downsample_point_cloud_factor: int | None = 16 - """Downsample the point cloud generated from each image by a factor. - - Useful to reduce memory usage when generating a large number of viewpoints. - Moreover, for high resolution cameras, the point cloud can be very dense without points providing additional information. + will be split into slices and for each slice, the number of sample points as given to the + :func:`nav_suite.collectors.sensor_data_sampling.SensorDataSampling.sample_sensor_data` method is generated. """ - downsample_point_cloud_voxel_size: float | None = 0.05 - """Voxel size to use for the downsampling of the point cloud. If None, no downsampling is applied.""" - slice_pc: bool = False - """Whether to slice the point cloud into the same slices as the viewpoints.""" ### # Saving Configuration ### save_path: str | None = None - """Directory to save the viewpoint samples, camera intrinsics and rendered images to. + """Directory to save the sensor data samples. If None, the directory is the same as the one of the obj file. Default is None.""" diff --git a/exts/nav_suite/nav_suite/collectors/sensors/__init__.py b/exts/nav_suite/nav_suite/collectors/sensors/__init__.py new file mode 100644 index 0000000..18bd27b --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensors/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from .base import SensorBase +from .base_cfg import SensorBaseCfg +from .camera import CameraSensor +from .camera_cfg import CameraSensorCfg +from .raycaster import RaycasterSensor +from .raycaster_cfg import RaycasterSensorCfg diff --git a/exts/nav_suite/nav_suite/collectors/sensors/base.py b/exts/nav_suite/nav_suite/collectors/sensors/base.py new file mode 100644 index 0000000..13c8546 --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensors/base.py @@ -0,0 +1,61 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from isaaclab.scene import InteractiveScene + +if TYPE_CHECKING: + from .base_cfg import SensorBaseCfg + + +class SensorBase(ABC): + """Base class for data extraction logic of sensor data. + + The individual function are executed in the following logic: + + pre_collection + + for i in range(num_collection_rounds): + pre_sim_update + scene.write_data_to_sim() + post_sim_update + + post_collection + """ + + def __init__(self, cfg: SensorBaseCfg, scene: InteractiveScene): + self.cfg = cfg + self.scene = scene + + @abstractmethod + def pre_collection(self, samples: torch.Tensor, filedir: str): + """Pre collection hook.""" + pass + + @abstractmethod + def post_collection(self, samples: torch.Tensor, filedir: str): + """Post collection hook.""" + pass + + @abstractmethod + def pre_sim_update(self, positions: torch.Tensor, orientations: torch.Tensor, env_ids: torch.Tensor): + """Pre simulation step hook.""" + pass + + @abstractmethod + def post_sim_update( + self, + samples_idx: torch.Tensor, + env_ids: torch.Tensor, + filedir: str, + slice_bounding_box: tuple[float, float, float, float] | None = None, + ): + """Post simulation step hook.""" + pass diff --git a/exts/nav_suite/nav_suite/collectors/sensors/base_cfg.py b/exts/nav_suite/nav_suite/collectors/sensors/base_cfg.py new file mode 100644 index 0000000..05dea89 --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensors/base_cfg.py @@ -0,0 +1,23 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.utils.configclass import configclass + +from .base import SensorBase + + +@configclass +class SensorBaseCfg: + """Base configuration for sensors.""" + + class_type: type[SensorBase] = SensorBase + """Class type of the sensor.""" + + requires_render: bool = False + """Whether the sensor requires render steps to fill the buffers. + + For USD/ tiled cameras, this has to be set to True. For raycaster cameras, this should be set to False to speedup + the sampling process. + """ diff --git a/exts/nav_suite/nav_suite/collectors/sensors/camera.py b/exts/nav_suite/nav_suite/collectors/sensors/camera.py new file mode 100644 index 0000000..5d2fc6d --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensors/camera.py @@ -0,0 +1,269 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import cv2 +import numpy as np +import open3d as o3d +import os +import time +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +import omni.log +from isaaclab.scene import InteractiveScene +from isaaclab.sensors import Camera, RayCasterCamera +from isaaclab.sim import SimulationContext + +from .base import SensorBase + +if TYPE_CHECKING: + from .camera_cfg import CameraSensorCfg + + +class CameraSensor(SensorBase): + """Camera sensor. + + The resulting folder structure is as follows: + + ``` graphql + cfg.data_dir + ├── camera_poses.txt # format: x y z qw qx qy qz + ├── cfg.cameras.keys[0] # required + | ├── intrinsics.txt # K-matrix (3x3) + | ├── annotator_0 # first annotator in the dict of the first camera + | | ├── xxxx.png # images saved with 4 digits, e.g. 0000.png + | ├── annotator_1 # second annotator in the dict of the first camera (optional) + | | ├── xxxx.png # images saved with 4 digits, e.g. 0000.png + ├── cfg.cameras.keys[1] # optional + | ├── intrinsics.txt # K-matrix (3x3) + | ├── annotator_0 # first annotator in the dict of the second camera + | | ├── xxxx.png # images saved with 4 digits, e.g. 0000.png + | ├── annotator_1 # second annotator in the dict of the second camera (optional) + | | ├── xxxx.png # images saved with 4 digits, e.g. 0000.png + ``` + + """ + + cfg: CameraSensorCfg + + def __init__(self, cfg: CameraSensorCfg, scene: InteractiveScene): + super().__init__(cfg, scene) + + # get sim context + self.sim = SimulationContext.instance() + + # init image index buffer + self._image_idx = [0] * len(self.cfg.cameras) + + # init time buffer + self._used_time = 0.0 + + # init buffer for point cloud + if self.cfg.generate_point_cloud: + self._pcd_points = [] + + def pre_collection(self, samples: torch.Tensor, filedir: str): + # create directory for images + filedir = os.path.join(filedir, "images") + for cam, annotator in self.cfg.cameras.items(): + [os.makedirs(os.path.join(filedir, cam, curr_annotator), exist_ok=True) for curr_annotator in annotator] + + # init a filter samples object + self._sample_filter = torch.zeros(samples.shape[0], dtype=torch.bool, device=self.scene.device) + + # save camera configurations + omni.log.info(f"Saving camera configurations to {filedir}.") + for cam in self.cfg.cameras.keys(): + # perform render updates to fill buffers if usd cameras are used + if isinstance(self.scene.sensors[cam], Camera): + for _ in range(5): + self.sim.render() + np.savetxt( + os.path.join(filedir, cam, "intrinsics.txt"), + self.scene.sensors[cam].data.intrinsic_matrices[0].cpu().numpy(), + delimiter=",", + ) + + def post_collection(self, samples: torch.Tensor, filedir: str): + # Save the complete point cloud as PLY + if self.cfg.generate_point_cloud: + pcd_points = np.concatenate(self._pcd_points, axis=0) + print(f"Generating point cloud from {len(pcd_points)} points.") + if self.cfg.downsample_point_cloud_voxel_size is not None: + pcd_points = ( + np.unique((pcd_points // self.cfg.downsample_point_cloud_voxel_size).astype(np.int32), axis=0) + * self.cfg.downsample_point_cloud_voxel_size + ) + + ply_filename = os.path.join(filedir, "images", "point_cloud.ply") + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pcd_points) + assert o3d.io.write_point_cloud(ply_filename, pcd) + omni.log.info(f"Saved complete point cloud to {ply_filename}.") + + # save camera poses + # TODO: check that this updates the samples in the SensorDataSampling object + samples = samples[~self._sample_filter] + np.savetxt(os.path.join(filedir, "images", "camera_poses.txt"), samples.cpu().numpy(), delimiter=",") + + ### + # Executed logic around the simulation update + # + # pre_sim_update + # scene.write_data_to_sim() + # post_sim_update + # + ### + + def pre_sim_update(self, positions: torch.Tensor, orientations: torch.Tensor, env_ids: torch.Tensor): + # set camera positions + for cam in self.cfg.cameras.keys(): + self.scene.sensors[cam].set_world_poses( + positions=positions, + orientations=orientations, + env_ids=env_ids, + convention="world", + ) + + def post_sim_update( + self, + samples_idx: torch.Tensor, + env_ids: torch.Tensor, + filedir: str, + slice_bounding_box: tuple[float, float, float, float] | None = None, + ): + # start time + start_time = time.time() + + # collect the output + image_data = {} + for cam_idx, curr_cam_annotator in enumerate(self.cfg.cameras.items()): + cam, annotator = curr_cam_annotator + + for curr_annotator in annotator: + image_data[(cam, curr_annotator)] = self.scene.sensors[cam].data.output[curr_annotator] + + # check validaty of the image data + image_filter = self.check_image_validity(image_data) + self._sample_filter[samples_idx] = image_filter[env_ids] + + # image modifiers + image_data = self.modify_images(image_data) + + # render + for cam_idx, curr_cam_annotator in enumerate(self.cfg.cameras.items()): + cam, annotator = curr_cam_annotator + + # save images + for idx in range(samples_idx.shape[0]): + if image_filter[idx]: + continue + + for curr_annotator in annotator: + # semantic segmentation or RGB + if ( + image_data[(cam, curr_annotator)].shape[-1] == 3 + or image_data[(cam, curr_annotator)].shape[-1] == 4 + ): + assert cv2.imwrite( + os.path.join( + filedir, "images", cam, curr_annotator, f"{self._image_idx[cam_idx]}".zfill(4) + ".png" + ), + cv2.cvtColor( + image_data[(cam, curr_annotator)][idx].cpu().numpy().astype(np.uint8), + cv2.COLOR_RGB2BGR, + ), + ) + # depth + else: + depth_image = image_data[(cam, curr_annotator)][idx] + assert cv2.imwrite( + os.path.join( + filedir, "images", cam, curr_annotator, f"{self._image_idx[cam_idx]}".zfill(4) + ".png" + ), + np.uint16(depth_image.cpu().numpy() * self.cfg.depth_scale), + ) + + # Check if point cloud generation is enabled + if self.cfg.generate_point_cloud: + # Convert depth image to point cloud + intrinsics = self.scene.sensors[cam].data.intrinsic_matrices[0] + points = math_utils.unproject_depth( + depth_image, intrinsics, is_ortho=annotator == "distance_to_image_plane" + ) + points = points.reshape(-1, 3) + + # transform points to world frame + points = math_utils.transform_points( + points, + self.scene.sensors[cam].data.pos_w[idx], + self.scene.sensors[cam].data.quat_w_ros[idx], + ) + + # filter points that are clipped + if self.scene.sensors[cam].cfg.depth_clipping_behavior == "zero": + point_filter = depth_image > 0 + elif self.scene.sensors[cam].cfg.depth_clipping_behavior == "max" and isinstance( + self.scene.sensors[cam], Camera + ): + point_filter = depth_image < self.scene.sensors[cam].cfg.spawn.clipping_range[1] + elif self.scene.sensors[cam].cfg.depth_clipping_behavior == "max" and isinstance( + self.scene.sensors[cam], RayCasterCamera + ): + point_filter = depth_image < self.scene.sensors[cam].cfg.max_distance + else: + point_filter = torch.ones_like(depth_image, dtype=torch.bool) + + points = points[point_filter.transpose(0, 1).reshape(-1)] + + # filter points that are outside the slice bounding box + # bounding box is in the format [x_max, y_max, x_min, y_min] + if slice_bounding_box is not None and self.cfg.slice_pc: + point_filter = ( + (points[:, 0] < slice_bounding_box[0]) + & (points[:, 0] > slice_bounding_box[2]) + & (points[:, 1] < slice_bounding_box[1]) + & (points[:, 1] > slice_bounding_box[3]) + ) + points = points[point_filter] + + if self.cfg.downsample_point_cloud_factor is not None: + points = points[:: self.cfg.downsample_point_cloud_factor] + if self.cfg.downsample_point_cloud_voxel_size is not None: + points = ( + torch.unique( + (points // self.cfg.downsample_point_cloud_voxel_size).type(torch.int32), + dim=0, + ) + * self.cfg.downsample_point_cloud_voxel_size + ) + + # Append points to the complete point cloud + self._pcd_points.append(points.cpu().numpy()) + + self._image_idx[cam_idx] += 1 + + if sum(self._image_idx) % 100 == 0: + print( + f"Rendered {sum(self._image_idx)} images in" + f" {(time.time() - start_time + self._used_time):.4f}s." + ) + + # update time + self._used_time = time.time() - start_time + + ### + # Helper functions + ### + + def check_image_validity(self, image_data: dict) -> torch.Tensor: + """Return False for valid images.""" + return torch.zeros(self.scene.num_envs, dtype=torch.bool, device=self.scene.device) + + def modify_images(self, image_data: dict) -> dict: + return image_data diff --git a/exts/nav_suite/nav_suite/collectors/sensors/camera_cfg.py b/exts/nav_suite/nav_suite/collectors/sensors/camera_cfg.py new file mode 100644 index 0000000..4a71b2d --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensors/camera_cfg.py @@ -0,0 +1,48 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.utils.configclass import configclass + +from .base_cfg import SensorBaseCfg +from .camera import CameraSensor + + +@configclass +class CameraSensorCfg(SensorBaseCfg): + """Camera sensor configuration.""" + + class_type: type[CameraSensor] = CameraSensor + """Class type of the sensor.""" + + ### + # Camera Configuration + ### + cameras: dict[str, list[str]] = { + "camera_0": ["semantic_segmentation"], + "camera_1": ["distance_to_image_plane"], + } + """Dict of cameras and corresponding annotators to use for the sensor data sampling.""" + + depth_scale: float = 1000.0 + """Scaling factor for the depth values.""" + + ### + # Point Cloud Generation + ### + generate_point_cloud: bool = True + """Whether to generate a point cloud from the depth images.""" + + downsample_point_cloud_factor: int | None = 16 + """Downsample the point cloud generated from each image by a factor. + + Useful to reduce memory usage when generating a large number of data samples. + Moreover, for high resolution cameras, the point cloud can be very dense without points providing additional information. + """ + + downsample_point_cloud_voxel_size: float | None = 0.05 + """Voxel size to use for the downsampling of the point cloud. If None, no downsampling is applied.""" + + slice_pc: bool = False + """Whether to slice the point cloud into the same slices as the image viewpoints.""" diff --git a/exts/nav_suite/nav_suite/collectors/sensors/raycaster.py b/exts/nav_suite/nav_suite/collectors/sensors/raycaster.py new file mode 100644 index 0000000..632b59d --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensors/raycaster.py @@ -0,0 +1,129 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import numpy as np +import os +import torch +from typing import TYPE_CHECKING + +from isaaclab.scene import InteractiveScene +from isaaclab.sensors import RayCaster +from isaaclab.utils import noise + +from .base import SensorBase + +if TYPE_CHECKING: + from .raycaster_cfg import RaycasterSensorCfg + + +class RaycasterSensor(SensorBase): + """Raycaster sensor.""" + + cfg: RaycasterSensorCfg + + def __init__(self, cfg: RaycasterSensorCfg, scene: InteractiveScene): + super().__init__(cfg, scene) + + # init index buffer + self._idx = 0 + + # check that the sensor is a raycaster + if not isinstance(self.scene.sensors[self.cfg.sensor_name], RayCaster): + raise ValueError("Sensor is not a raycaster.") + + # check that the observations term cfg is a list + if self.cfg.obs_term_cfg is not None and not isinstance(self.cfg.obs_term_cfg, list): + self.cfg.obs_term_cfg = [self.cfg.obs_term_cfg] + + def pre_collection(self, samples: torch.Tensor, filedir: str): + # create directory for raycaster data + filedir = os.path.join(filedir, "raycaster") + os.makedirs(filedir, exist_ok=True) + + def post_collection(self, samples: torch.Tensor, filedir: str): + # save raycaster poses + np.savetxt(os.path.join(filedir, "raycaster_poses.txt"), samples.cpu().numpy(), delimiter=",") + + def pre_sim_update(self, positions: torch.Tensor, orientations: torch.Tensor, env_ids: torch.Tensor): + # as the raycaster itself is a virtual sensor, it does not directly support to set the world poses + # instead, we need to set the world poses of the xform / articulation / rigid body under which it is spawned + + try: + asset = self.scene.articulations[self.cfg.asset_name] + except KeyError: + asset = None + + if asset is None: + try: + asset = self.scene.rigid_objects[self.cfg.asset_name] + except KeyError: + asset = None + + assert asset is not None, "Asset name not found in scene. Has to be an articulation or rigid body." + + # set the world pose of the asset + asset.write_root_pose_to_sim( + torch.cat([positions[env_ids], orientations[env_ids]], dim=-1), + env_ids=env_ids.to(self.scene.device), + ) + + def post_sim_update( + self, + samples_idx: torch.Tensor, + env_ids: torch.Tensor, + filedir: str, + slice_bounding_box: tuple[float, float, float, float] | None = None, + ): + if self.cfg.obs_term_cfg is not None: + # construct a dummy class with the current scene (no other attribute can be accesses) + + class ManagerBasedEnvDummy: + def __init__(self, scene: InteractiveScene): + self.scene = scene + + env_dummy = ManagerBasedEnvDummy(self.scene) + + output = {} + + for obs_term_cfg in self.cfg.obs_term_cfg: + # get the mdp output + curr_output = obs_term_cfg.func(env_dummy, **obs_term_cfg.params) + + # apply post-processing + if obs_term_cfg.modifiers is not None: + for modifier in obs_term_cfg.modifiers: + curr_output = modifier.func(curr_output, **modifier.params) + if isinstance(obs_term_cfg.noise, noise.NoiseCfg): + curr_output = obs_term_cfg.noise.func(curr_output, obs_term_cfg.noise) + elif isinstance(obs_term_cfg.noise, noise.NoiseModelCfg) and obs_term_cfg.noise.func is not None: + curr_output = obs_term_cfg.noise.func(curr_output) + if obs_term_cfg.clip: + curr_output = curr_output.clip_(min=obs_term_cfg.clip[0], max=obs_term_cfg.clip[1]) + if obs_term_cfg.scale is not None: + curr_output = curr_output.mul_(obs_term_cfg.scale) + + # save the output for the current function + output[obs_term_cfg.func.__name__] = curr_output + + else: + output = self.scene.sensors[self.cfg.sensor_name].data.ray_hits_w + + for env_id in env_ids: + if isinstance(output, dict): + for key, value in output.items(): + np.save( + os.path.join(filedir, "raycaster", key, f"{self._idx}".zfill(4) + ".npy"), + value[env_id].cpu().numpy(), + ) + else: + # save the output + np.save( + os.path.join(filedir, "raycaster", f"{self._idx}".zfill(4) + ".npy"), output[env_id].cpu().numpy() + ) + + # increment the index + self._idx += 1 diff --git a/exts/nav_suite/nav_suite/collectors/sensors/raycaster_cfg.py b/exts/nav_suite/nav_suite/collectors/sensors/raycaster_cfg.py new file mode 100644 index 0000000..3292230 --- /dev/null +++ b/exts/nav_suite/nav_suite/collectors/sensors/raycaster_cfg.py @@ -0,0 +1,40 @@ +# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import MISSING + +from isaaclab.managers import ObservationTermCfg +from isaaclab.utils.configclass import configclass + +from .base_cfg import SensorBaseCfg +from .raycaster import RaycasterSensor + + +@configclass +class RaycasterSensorCfg(SensorBaseCfg): + """Raycaster sensor configuration.""" + + class_type: type[RaycasterSensor] = RaycasterSensor + """Class type of the sensor.""" + + ### + # Raycaster Configuration + ### + + sensor_name: str = MISSING + """Name of the sensor to use for the raycaster.""" + + asset_name: str = MISSING + """Name of the asset under which the raycaster is spawned. + + Need to access the asset to set the world pose of the raycaster as the sensor itself is virtual and therefore does + not support to set the world pose directly. + """ + + obs_term_cfg: list[ObservationTermCfg] | ObservationTermCfg | None = None + """Observation term configuration. + + If provided, the output of the observation term will be saved. If None, the :attr:`ray_hits_w` will be saved. + """ diff --git a/exts/nav_suite/nav_suite/collectors/viewpoint_sampling.py b/exts/nav_suite/nav_suite/collectors/viewpoint_sampling.py deleted file mode 100644 index a5a5172..0000000 --- a/exts/nav_suite/nav_suite/collectors/viewpoint_sampling.py +++ /dev/null @@ -1,483 +0,0 @@ -# Copyright (c) 2025, The Nav-Suite Project Developers (https://github.com/leggedrobotics/nav-suite/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import cv2 -import numpy as np -import open3d as o3d -import os -import pickle -import random -import time -import torch - -import isaaclab.utils.math as math_utils -import omni.log -from isaaclab.markers import VisualizationMarkers -from isaaclab.markers.config import GREEN_ARROW_X_MARKER_CFG -from isaaclab.scene import InteractiveScene -from isaaclab.sensors import Camera, RayCasterCamera -from isaaclab.sim import SimulationContext - -from .viewpoint_sampling_cfg import ViewpointSamplingCfg - - -class ViewpointSampling: - def __init__(self, cfg: ViewpointSamplingCfg, scene: InteractiveScene): - # save cfg and env - self.cfg = cfg - self.scene = scene - - # get sim context - self.sim = SimulationContext.instance() - - # analyse terrains -- check if singleton is used and available - if ( - hasattr(self.cfg.terrain_analysis.class_type, "instance") - and self.cfg.terrain_analysis.class_type.instance() is not None - ): - self.terrain_analyser = self.cfg.terrain_analysis.class_type.instance() - else: - self.terrain_analyser = self.cfg.terrain_analysis.class_type(self.cfg.terrain_analysis, scene=self.scene) - - ### - # Properties - ### - - @property - def samples(self) -> torch.Tensor | list[torch.Tensor]: - """Get the sampled viewpoints. - - Samples are stored in a torch tensor with the structure - [x, y, z, qw, qx, qv, qz] - """ - return self._samples - - @property - def sliced_bounding_boxes(self) -> list[tuple[float, float, float, float]] | None: - """Get the sliced bounding boxes. - - The sliced bounding boxes follow the terrain analysis format of the mesh dimensions, i.e. [x_max, y_max, x_min, y_min] - """ - if self.cfg.sliced_sampling is not None and hasattr(self, "_sliced_bounding_boxes"): - return self._sliced_bounding_boxes - else: - return None - - ### - # Operations - ### - - def sample_viewpoints(self, nbr_viewpoints: int, seed: int = 1) -> torch.Tensor | list[torch.Tensor]: - """Sample viewpoints for the given number of viewpoints and seed. - - Samples are stored in a torch tensor with the structure - [x, y, z, qw, qx, qv, qz] - - Args: - nbr_viewpoints (int): The number of viewpoints to sample. - seed (int, optional): The seed for the random number generator. Defaults to 1. - - Returns: - torch.Tensor | list[torch.Tensor]: The sampled viewpoints. - """ - - # get the number of slices and their locations - # NOTE: the number of viewpoints is kept constant per slice - if self.cfg.sliced_sampling is not None: - # get mesh size from terrain analysis -- [x_max, y_max, x_min, y_min] - mesh_size = self.terrain_analyser.mesh_dimensions - # get the number of slices in both x and y direction - nbr_slices_x = int(np.ceil((mesh_size[0] - mesh_size[2]) / self.cfg.sliced_sampling[0])) - nbr_slices_y = int(np.ceil((mesh_size[1] - mesh_size[3]) / self.cfg.sliced_sampling[1])) - slice_locations_x = np.linspace(mesh_size[2], mesh_size[0], nbr_slices_x) - slice_locations_y = np.linspace(mesh_size[3], mesh_size[1], nbr_slices_y) - # get the slice bounding boxes - # slide bounding boxes follow the terrain analysis format of the mesh dimensions, i.e. [x_max, y_max, x_min, y_min] - self._sliced_bounding_boxes = [] - for i in range(nbr_slices_x - 1): - for j in range(nbr_slices_y - 1): - self._sliced_bounding_boxes.append( - (slice_locations_x[i + 1], slice_locations_y[j + 1], slice_locations_x[i], slice_locations_y[j]) - ) - - # execute viewpoint sampling for each slice - self._samples = [] - failed_slices = [] - for i, slice_bounding_box in enumerate(self._sliced_bounding_boxes): - try: - self._samples.append(self._sample_viewpoint_per_area(nbr_viewpoints, seed, slice_bounding_box)) - except Exception as e: - failed_slices.append(slice_bounding_box) - omni.log.warn(f"Error sampling viewpoints for slice {slice_bounding_box}: {e}") - continue - - # remove failed slices - self._sliced_bounding_boxes = [box for box in self._sliced_bounding_boxes if box not in failed_slices] - # dump the sliced bounding boxes and samples - filedir = self.cfg.save_path if self.cfg.save_path else self._get_save_filedir() - with open(os.path.join(filedir, "sliced_bounding_boxes.pkl"), "wb") as f: - pickle.dump(self._sliced_bounding_boxes, f) - - else: - # execute viewpoint sampling for the whole mesh - self._samples = self._sample_viewpoint_per_area(nbr_viewpoints, seed) - - return self._samples - - def render_viewpoints( - self, samples: torch.Tensor | list[torch.Tensor] | None = None - ) -> torch.Tensor | list[torch.Tensor]: - """Render the images at the given viewpoints and save them to the drive.""" - if samples is None: - samples = self.samples - - if isinstance(samples, list): - assert ( - self.sliced_bounding_boxes is not None - ), "Sliced bounding boxes must be set to render a list of viewpoints samples." - assert len(samples) == len( - self.sliced_bounding_boxes - ), "The number of samples must match the number of sliced bounding boxes." - render_sample_points = [] - for sample, slice_bounding_box in zip(samples, self.sliced_bounding_boxes): - print(f"[INFO] Rendering viewpoints for slice {slice_bounding_box}") - render_sample_points.append(self._render_viewpoints_per_area(sample, slice_bounding_box)) - else: - render_sample_points = self._render_viewpoints_per_area(samples) - - return render_sample_points - - def check_image_validity(self, image_data: dict) -> torch.Tensor: - """Return False for valid images.""" - return torch.zeros(self.scene.num_envs, dtype=torch.bool, device=self.scene.device) - - def modify_images(self, image_data: dict) -> dict: - return image_data - - ### - # Helper functions - ### - - def _sample_viewpoint_per_area( - self, nbr_viewpoints: int, seed: int = 1, slice_bounding_box: tuple[float, float, float, float] | None = None - ) -> torch.Tensor: - - filedir = self.cfg.save_path if self.cfg.save_path else self._get_save_filedir() - if slice_bounding_box is not None: - slice_bounding_box_str = "slice_" + "_".join( - f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in slice_bounding_box - ) - filedir = os.path.join(filedir, slice_bounding_box_str) - elif self.cfg.terrain_analysis.terrain_bounding_box is not None: - slice_bounding_box_str = "slice_" + "_".join( - f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in self.cfg.terrain_analysis.terrain_bounding_box - ) - filedir = os.path.join(filedir, slice_bounding_box_str) - else: - slice_bounding_box_str = "full" - filename = os.path.join(filedir, f"viewpoints_seed{seed}_samples{nbr_viewpoints}.pkl") - if os.path.isfile(filename): - with open(filename, "rb") as f: - data = pickle.load(f) - # add loaded path dict to data dict - omni.log.info(f"Loaded {nbr_viewpoints} with seed {seed} for map part {slice_bounding_box_str}.") - return data - else: - omni.log.info( - f"No viewpoint samples found for seed {seed} and {nbr_viewpoints} samples for map part" - f" {slice_bounding_box_str}." - ) - - print(f"[INFO] Sampling viewpoints for {slice_bounding_box_str}") - - # analyse terrain if not done yet - if slice_bounding_box is not None: - # overwrite the mesh dimensions to influence where the sampling is done - self.terrain_analyser.reset_graph() - self.terrain_analyser._mesh_dimensions = slice_bounding_box - self.terrain_analyser.analyse() - elif not self.terrain_analyser.complete: - self.terrain_analyser.analyse() - - # set seed - random.seed(seed) - omni.log.info(f"Start sampling {nbr_viewpoints} viewpoints for map part {slice_bounding_box_str}.") - - # samples are organized in [point_idx, neighbor_idx, distance] - # sample from each point the neighbor with the largest distance - nbr_samples_per_point = int(np.ceil(nbr_viewpoints / self.terrain_analyser.points.shape[0]).item()) - sample_locations = torch.zeros((nbr_samples_per_point * self.terrain_analyser.points.shape[0], 2)) - sample_locations_count = 0 - curr_point_idx = 0 - while sample_locations_count < nbr_viewpoints: - # get samples - sample_idx = self.terrain_analyser.samples[:, 0] == curr_point_idx - sample_idx_select = torch.randperm(sample_idx.sum())[:nbr_samples_per_point] - sample_locations[sample_locations_count : sample_locations_count + sample_idx_select.shape[0]] = ( - self.terrain_analyser.samples[sample_idx][sample_idx_select, :2] - ) - sample_locations_count += sample_idx_select.shape[0] - curr_point_idx += 1 - # reset point index if all points are sampled - if curr_point_idx >= self.terrain_analyser.points.shape[0]: - curr_point_idx = 0 - - sample_locations = sample_locations[:sample_locations_count].type(torch.int64) - - # get the z angle of the neighbor that is closest to the origin point - neighbor_direction = ( - self.terrain_analyser.points[sample_locations[:, 0]] - self.terrain_analyser.points[sample_locations[:, 1]] - ) - z_angles = torch.atan2(neighbor_direction[:, 1], neighbor_direction[:, 0]).to("cpu") - - # vary the rotation of the forward and horizontal axis (in camera frame) as a uniform distribution within the limits - x_angles = math_utils.sample_uniform( - self.cfg.x_angle_range[0], self.cfg.x_angle_range[1], sample_locations_count, device="cpu" - ) - y_angles = math_utils.sample_uniform( - self.cfg.y_angle_range[0], self.cfg.y_angle_range[1], sample_locations_count, device="cpu" - ) - x_angles = torch.deg2rad(x_angles) - y_angles = torch.deg2rad(y_angles) - - samples = torch.zeros((sample_locations_count, 7)) - samples[:, :3] = self.terrain_analyser.points[sample_locations[:, 0]] - samples[:, 3:] = math_utils.quat_from_euler_xyz(x_angles, y_angles, z_angles) - - omni.log.info(f"Sampled {sample_locations_count} viewpoints for map part {slice_bounding_box_str}.") - - # save samples - os.makedirs(filedir, exist_ok=True) - with open(filename, "wb") as f: - pickle.dump(samples, f) - - omni.log.info(f"Saved {sample_locations_count} viewpoints with seed {seed} to {filename}.") - - # debug points and orientation - if self.cfg.debug_viz: - env_render_steps = 1000 - marker_cfg = GREEN_ARROW_X_MARKER_CFG.copy() - marker_cfg.prim_path = "/Visuals/viewpoints" - marker_cfg.markers["arrow"].scale = (0.1, 0.1, 0.1) - self.visualizer = VisualizationMarkers(marker_cfg) - self.visualizer.visualize(samples[:, :3], samples[:, 3:]) - - omni.log.info(f"Visualizing {sample_locations_count} samples for {env_render_steps} render steps...") - for _ in range(env_render_steps): - self.sim.render() - - self.visualizer.set_visibility(False) - omni.log.info("Done visualizing.") - - return samples - - def _render_viewpoints_per_area( - self, samples: torch.Tensor, slice_bounding_box: tuple[float, float, float, float] | None = None - ) -> torch.Tensor: - omni.log.info(f"Start rendering {samples.shape[0]} images.") - # define how many rounds are necessary to render all viewpoints - num_rounds = int(np.ceil(samples.shape[0] / self.scene.num_envs)) - # image_idx - image_idx = [0] * len(self.cfg.cameras) - - filedir = self.cfg.save_path if self.cfg.save_path else self._get_save_filedir() - if slice_bounding_box is not None: - slice_bounding_box_str = "slice_" + "_".join( - f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in slice_bounding_box - ) - filedir = os.path.join(filedir, slice_bounding_box_str) - elif self.cfg.terrain_analysis.terrain_bounding_box is not None: - slice_bounding_box_str = "slice_" + "_".join( - f"{'n' if x < 0 else ''}{abs(x):.1f}" for x in self.cfg.terrain_analysis.terrain_bounding_box - ) - filedir = os.path.join(filedir, slice_bounding_box_str) - filedir = os.path.join(filedir, "images") - for cam, annotator in self.cfg.cameras.items(): - [os.makedirs(os.path.join(filedir, cam, curr_annotator), exist_ok=True) for curr_annotator in annotator] - - # save camera configurations - omni.log.info(f"Saving camera configurations to {filedir}.") - for cam in self.cfg.cameras.keys(): - # perform render steps to fill buffers if usd cameras are used - if isinstance(self.scene.sensors[cam], Camera): - for _ in range(5): - self.sim.render() - np.savetxt( - os.path.join(filedir, cam, "intrinsics.txt"), - self.scene.sensors[cam].data.intrinsic_matrices[0].cpu().numpy(), - delimiter=",", - ) - - # init points array for point cloud - if self.cfg.generate_point_cloud: - pcd_points = [] - - # save images - samples = samples.to(self.scene.device) - sample_filter = torch.zeros(samples.shape[0], dtype=torch.bool, device=self.scene.device) - start_time = time.time() - for i in range(num_rounds): - # get samples idx - samples_idx = torch.arange(i * self.scene.num_envs, min((i + 1) * self.scene.num_envs, samples.shape[0])) - env_ids = torch.arange(samples_idx.shape[0]) - # set camera positions - for cam in self.cfg.cameras.keys(): - self.scene.sensors[cam].set_world_poses( - positions=samples[samples_idx, :3], - orientations=samples[samples_idx, 3:], - env_ids=env_ids, - convention="world", - ) - # update simulation - self.scene.write_data_to_sim() - # perform render steps to fill buffers if usd cameras are used - if any([isinstance(self.scene.sensors[cam], Camera) for cam in self.cfg.cameras.keys()]): - for _ in range(10): - self.sim.render() - # update scene buffers - self.scene.update(self.sim.get_physics_dt()) - - # collect the output - image_data = {} - for cam_idx, curr_cam_annotator in enumerate(self.cfg.cameras.items()): - cam, annotator = curr_cam_annotator - - for curr_annotator in annotator: - image_data[(cam, curr_annotator)] = self.scene.sensors[cam].data.output[curr_annotator] - - # check validaty of the image data - image_filter = self.check_image_validity(image_data) - sample_filter[samples_idx] = image_filter[env_ids] - - # image modifiers - image_data = self.modify_images(image_data) - - # render - for cam_idx, curr_cam_annotator in enumerate(self.cfg.cameras.items()): - cam, annotator = curr_cam_annotator - - # save images - for idx in range(samples_idx.shape[0]): - if image_filter[idx]: - continue - - for curr_annotator in annotator: - # semantic segmentation or RGB - if ( - image_data[(cam, curr_annotator)].shape[-1] == 3 - or image_data[(cam, curr_annotator)].shape[-1] == 4 - ): - assert cv2.imwrite( - os.path.join(filedir, cam, curr_annotator, f"{image_idx[cam_idx]}".zfill(4) + ".png"), - cv2.cvtColor( - image_data[(cam, curr_annotator)][idx].cpu().numpy().astype(np.uint8), - cv2.COLOR_RGB2BGR, - ), - ) - # depth - else: - depth_image = image_data[(cam, curr_annotator)][idx] - assert cv2.imwrite( - os.path.join(filedir, cam, curr_annotator, f"{image_idx[cam_idx]}".zfill(4) + ".png"), - np.uint16(depth_image.cpu().numpy() * self.cfg.depth_scale), - ) - - # Check if point cloud generation is enabled - if self.cfg.generate_point_cloud: - # Convert depth image to point cloud - intrinsics = self.scene.sensors[cam].data.intrinsic_matrices[0] - points = math_utils.unproject_depth( - depth_image, intrinsics, is_ortho=annotator == "distance_to_image_plane" - ) - points = points.reshape(-1, 3) - - # transform points to world frame - points = math_utils.transform_points( - points, - self.scene.sensors[cam].data.pos_w[idx], - self.scene.sensors[cam].data.quat_w_ros[idx], - ) - - # filter points that are clipped - if self.scene.sensors[cam].cfg.depth_clipping_behavior == "zero": - point_filter = depth_image > 0 - elif self.scene.sensors[cam].cfg.depth_clipping_behavior == "max" and isinstance( - self.scene.sensors[cam], Camera - ): - point_filter = depth_image < self.scene.sensors[cam].cfg.spawn.clipping_range[1] - elif self.scene.sensors[cam].cfg.depth_clipping_behavior == "max" and isinstance( - self.scene.sensors[cam], RayCasterCamera - ): - point_filter = depth_image < self.scene.sensors[cam].cfg.max_distance - else: - point_filter = torch.ones_like(depth_image, dtype=torch.bool) - - points = points[point_filter.transpose(0, 1).reshape(-1)] - - # filter points that are outside the slice bounding box - # bounding box is in the format [x_max, y_max, x_min, y_min] - if slice_bounding_box is not None and self.cfg.slice_pc: - point_filter = ( - (points[:, 0] < slice_bounding_box[0]) - & (points[:, 0] > slice_bounding_box[2]) - & (points[:, 1] < slice_bounding_box[1]) - & (points[:, 1] > slice_bounding_box[3]) - ) - points = points[point_filter] - - if self.cfg.downsample_point_cloud_factor is not None: - points = points[:: self.cfg.downsample_point_cloud_factor] - if self.cfg.downsample_point_cloud_voxel_size is not None: - points = ( - torch.unique( - (points // self.cfg.downsample_point_cloud_voxel_size).type(torch.int32), - dim=0, - ) - * self.cfg.downsample_point_cloud_voxel_size - ) - - # Append points to the complete point cloud - pcd_points.append(points.cpu().numpy()) - - image_idx[cam_idx] += 1 - - if sum(image_idx) % 100 == 0: - print(f"Rendered {sum(image_idx)} images in {(time.time() - start_time):.4f}s.") - - # Save the complete point cloud as PLY - if self.cfg.generate_point_cloud: - pcd_points = np.concatenate(pcd_points, axis=0) - print(f"Generating point cloud from {len(pcd_points)} points.") - if self.cfg.downsample_point_cloud_voxel_size is not None: - pcd_points = ( - np.unique((pcd_points // self.cfg.downsample_point_cloud_voxel_size).astype(np.int32), axis=0) - * self.cfg.downsample_point_cloud_voxel_size - ) - - ply_filename = os.path.join(filedir, "point_cloud.ply") - pcd = o3d.geometry.PointCloud() - pcd.points = o3d.utility.Vector3dVector(pcd_points) - assert o3d.io.write_point_cloud(ply_filename, pcd) - omni.log.info(f"Saved complete point cloud to {ply_filename}.") - - # save camera poses - samples = samples[~sample_filter] - np.savetxt(os.path.join(filedir, "camera_poses.txt"), samples.cpu().numpy(), delimiter=",") - - return samples - - def _get_save_filedir(self) -> str: - # get env name - if isinstance(self.scene.terrain.cfg.usd_path, str): - terrain_file_path = self.scene.terrain.cfg.usd_path - else: - raise KeyError("Only implemented for terrains loaded from usd and matterport") - env_name = os.path.splitext(terrain_file_path)[0] - # create directory if necessary - filedir = os.path.join(terrain_file_path, env_name) - os.makedirs(filedir, exist_ok=True) - return filedir diff --git a/scripts/nav_suite/collector/carla_viewpoint_sampling.py b/scripts/nav_suite/collector/carla_sensor_data_sampling.py similarity index 90% rename from scripts/nav_suite/collector/carla_viewpoint_sampling.py rename to scripts/nav_suite/collector/carla_sensor_data_sampling.py index 8f5b334..e805c3e 100644 --- a/scripts/nav_suite/collector/carla_viewpoint_sampling.py +++ b/scripts/nav_suite/collector/carla_sensor_data_sampling.py @@ -38,7 +38,7 @@ from isaaclab.utils.timer import Timer from nav_suite import NAVSUITE_DATA_DIR -from nav_suite.collectors import ViewpointSampling, ViewpointSamplingCfg +from nav_suite.collectors import CameraSensorCfg, SensorDataSampling, SensorDataSamplingCfg from nav_suite.terrains import NavTerrainImporterCfg """ @@ -111,7 +111,7 @@ def main(): sim.set_camera_view([130, -125, 30], [100, -130, 0.5]) # setup sampling config - cfg = ViewpointSamplingCfg() + cfg = SensorDataSamplingCfg(sensor_data_handlers=[CameraSensorCfg()]) # overwrite semantic cost mapping and adjust parameters based on larger map cfg.terrain_analysis.semantic_cost_mapping = os.path.join( NAVSUITE_DATA_DIR, "unreal", "town01", "semantic_costs.yaml" @@ -132,15 +132,15 @@ def main(): with Timer("[INFO]: Time taken for simulation start", "simulation_start"): sim.reset() - explorer = ViewpointSampling(cfg, scene) + data_sampler = SensorDataSampling(cfg, scene) # Now we are ready! omni.log.info("Setup complete...") - # sample and render viewpoints - samples = explorer.sample_viewpoints(9560) - explorer.render_viewpoints(samples) + # sample and render sensor data + samples = data_sampler.sample_sensor_data(9560) + data_sampler.render_sensor_data(samples) print( - "Viewpoints sampled and rendered will continue to render the environment and visualize the last camera" + "Sensor data sampled and rendered will continue to render the environment and visualize the last camera" " positions..." ) @@ -151,7 +151,7 @@ def main(): # Perform step sim.render() # Update buffers - explorer.scene.update(sim_dt) + data_sampler.scene.update(sim_dt) if __name__ == "__main__": diff --git a/scripts/nav_suite/collector/matterport_viewpoint_sampling.py b/scripts/nav_suite/collector/matterport_sensor_data_sampling.py similarity index 82% rename from scripts/nav_suite/collector/matterport_viewpoint_sampling.py rename to scripts/nav_suite/collector/matterport_sensor_data_sampling.py index ae1b46e..bac9a2d 100644 --- a/scripts/nav_suite/collector/matterport_viewpoint_sampling.py +++ b/scripts/nav_suite/collector/matterport_sensor_data_sampling.py @@ -37,7 +37,7 @@ from isaaclab.utils.timer import Timer from nav_suite import NAVSUITE_DATA_DIR -from nav_suite.collectors import ViewpointSampling, ViewpointSamplingCfg +from nav_suite.collectors import CameraSensorCfg, SensorDataSampling, SensorDataSamplingCfg from nav_suite.terrain_analysis import TerrainAnalysisCfg """ @@ -53,11 +53,12 @@ def main(): sim_cfg = sim_utils.SimulationCfg() sim = SimulationContext(sim_cfg) - cfg = ViewpointSamplingCfg( + cfg = SensorDataSamplingCfg( + sensor_data_handlers=[CameraSensorCfg()], terrain_analysis=TerrainAnalysisCfg( semantic_cost_mapping=os.path.join(NAVSUITE_DATA_DIR, "matterport", "semantic_costs.yaml"), raycaster_sensor="camera_0", - ) + ), ) # construct the scene @@ -69,15 +70,15 @@ def main(): with Timer("[INFO]: Time taken for simulation start", "simulation_start"): sim.reset() - explorer = ViewpointSampling(cfg, scene) + data_sampler = SensorDataSampling(cfg, scene) # Now we are ready! omni.log.info("Setup complete...") - # sample and render viewpoints - samples = explorer.sample_viewpoints(args_cli.num_samples) - explorer.render_viewpoints(samples) + # sample and render sensor data + samples = data_sampler.sample_sensor_data(args_cli.num_samples) + data_sampler.render_sensor_data(samples) print( - "Viewpoints sampled and rendered will continue to render the environment and visualize the last camera" + "Sensor data sampled and rendered will continue to render the environment and visualize the last camera" " positions..." ) @@ -88,7 +89,7 @@ def main(): # Perform step sim.render() # Update buffers - explorer.scene.update(sim_dt) + data_sampler.scene.update(sim_dt) if __name__ == "__main__": diff --git a/scripts/nav_suite/collector/matterport_trajectory_sampling.py b/scripts/nav_suite/collector/matterport_trajectory_sampling.py index 2382fb3..1123f0b 100644 --- a/scripts/nav_suite/collector/matterport_trajectory_sampling.py +++ b/scripts/nav_suite/collector/matterport_trajectory_sampling.py @@ -75,7 +75,7 @@ def main(): # Now we are ready! omni.log.info("Setup complete...") - # sample viewpoints + # sample paths explorer.sample_paths(1000, 0.5, 10.0) print("Trajectories sampled and simulation will continue to render the environment...") diff --git a/scripts/tools/cut_usd_to_bounding_box.py b/scripts/tools/cut_usd_to_bounding_box.py index 6d494b9..7f20802 100644 --- a/scripts/tools/cut_usd_to_bounding_box.py +++ b/scripts/tools/cut_usd_to_bounding_box.py @@ -4,7 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 """ -Script to cut a USD file to a certain bounding box. +This script implements a tool to cut a USD scene to a bounding box. """ """Launch Isaac Sim Simulator first."""