# Copyright (c) 2022-2024, The ORBIT Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations
import torch
from collections.abc import Sequence
from tensordict import TensorDict
from typing import TYPE_CHECKING, ClassVar, Literal
import omni.physics.tensors.impl.api as physx
from omni.isaac.core.prims import XFormPrimView
import omni.isaac.orbit.utils.math as math_utils
from omni.isaac.orbit.sensors.camera import CameraData
from omni.isaac.orbit.sensors.camera.utils import convert_orientation_convention, create_rotation_matrix_from_view
from omni.isaac.orbit.utils.warp import raycast_mesh
from .ray_caster import RayCaster
if TYPE_CHECKING:
from .ray_caster_camera_cfg import RayCasterCameraCfg
[docs]class RayCasterCamera(RayCaster):
"""A ray-casting camera sensor.
The ray-caster camera uses a set of rays to get the distances to meshes in the scene. The rays are
defined in the sensor's local coordinate frame. The sensor has the same interface as the
:class:`omni.isaac.orbit.sensors.Camera` that implements the camera class through USD camera prims.
However, this class provides a faster image generation. The sensor converts meshes from the list of
primitive paths provided in the configuration to Warp meshes. The camera then ray-casts against these
Warp meshes only.
Currently, only the following annotators are supported:
- ``"distance_to_camera"``: An image containing the distance to camera optical center.
- ``"distance_to_image_plane"``: An image containing distances of 3D points from camera plane along camera's z-axis.
- ``"normals"``: An image containing the local surface normal vectors at each pixel.
.. note::
Currently, only static meshes are supported. Extending the warp mesh to support dynamic meshes
is a work in progress.
"""
cfg: RayCasterCameraCfg
"""The configuration parameters."""
UNSUPPORTED_TYPES: ClassVar[set[str]] = {
"rgb",
"instance_id_segmentation",
"instance_id_segmentation_fast",
"instance_segmentation",
"instance_segmentation_fast",
"semantic_segmentation",
"skeleton_data",
"motion_vectors",
"bounding_box_2d_tight",
"bounding_box_2d_tight_fast",
"bounding_box_2d_loose",
"bounding_box_2d_loose_fast",
"bounding_box_3d",
"bounding_box_3d_fast",
}
"""A set of sensor types that are not supported by the ray-caster camera."""
[docs] def __init__(self, cfg: RayCasterCameraCfg):
"""Initializes the camera object.
Args:
cfg: The configuration parameters.
Raises:
ValueError: If the provided data types are not supported by the ray-caster camera.
"""
# perform check on supported data types
self._check_supported_data_types(cfg)
# initialize base class
super().__init__(cfg)
# create empty variables for storing output data
self._data = CameraData()
def __str__(self) -> str:
"""Returns: A string containing information about the instance."""
return (
f"Ray-Caster-Camera @ '{self.cfg.prim_path}': \n"
f"\tview type : {self._view.__class__}\n"
f"\tupdate period (s) : {self.cfg.update_period}\n"
f"\tnumber of meshes : {len(RayCaster.meshes)}\n"
f"\tnumber of sensors : {self._view.count}\n"
f"\tnumber of rays/sensor: {self.num_rays}\n"
f"\ttotal number of rays : {self.num_rays * self._view.count}\n"
f"\timage shape : {self.image_shape}"
)
"""
Properties
"""
@property
def data(self) -> CameraData:
# update sensors if needed
self._update_outdated_buffers()
# return the data
return self._data
@property
def image_shape(self) -> tuple[int, int]:
"""A tuple containing (height, width) of the camera sensor."""
return (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width)
@property
def frame(self) -> torch.tensor:
"""Frame number when the measurement took place."""
return self._frame
"""
Operations.
"""
[docs] def set_intrinsic_matrices(
self, matrices: torch.Tensor, focal_length: float = 1.0, env_ids: Sequence[int] | None = None
):
"""Set the intrinsic matrix of the camera.
Args:
matrices: The intrinsic matrices for the camera. Shape is (N, 3, 3).
focal_length: Focal length to use when computing aperture values. Defaults to 1.0.
env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices.
"""
# resolve env_ids
if env_ids is None:
env_ids = slice(None)
# save new intrinsic matrices and focal length
self._data.intrinsic_matrices[env_ids] = matrices.to(self._device)
self._focal_length = focal_length
# recompute ray directions
self.ray_starts[env_ids], self.ray_directions[env_ids] = self.cfg.pattern_cfg.func(
self.cfg.pattern_cfg, self._data.intrinsic_matrices[env_ids], self._device
)
[docs] def reset(self, env_ids: Sequence[int] | None = None):
# reset the timestamps
super().reset(env_ids)
# resolve None
if env_ids is None:
env_ids = slice(None)
# reset the data
# note: this recomputation is useful if one performs events such as randomizations on the camera poses.
pos_w, quat_w = self._compute_camera_world_poses(env_ids)
self._data.pos_w[env_ids] = pos_w
self._data.quat_w_world[env_ids] = quat_w
# Reset the frame count
self._frame[env_ids] = 0
[docs] def set_world_poses(
self,
positions: torch.Tensor | None = None,
orientations: torch.Tensor | None = None,
env_ids: Sequence[int] | None = None,
convention: Literal["opengl", "ros", "world"] = "ros",
):
"""Set the pose of the camera w.r.t. the world frame using specified convention.
Since different fields use different conventions for camera orientations, the method allows users to
set the camera poses in the specified convention. Possible conventions are:
- :obj:`"opengl"` - forward axis: -Z - up axis +Y - Offset is applied in the OpenGL (Usd.Camera) convention
- :obj:`"ros"` - forward axis: +Z - up axis -Y - Offset is applied in the ROS convention
- :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention
See :meth:`omni.isaac.orbit.sensors.camera.utils.convert_orientation_convention` for more details
on the conventions.
Args:
positions: The cartesian coordinates (in meters). Shape is (N, 3).
Defaults to None, in which case the camera position in not changed.
orientations: The quaternion orientation in (w, x, y, z). Shape is (N, 4).
Defaults to None, in which case the camera orientation in not changed.
env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices.
convention: The convention in which the poses are fed. Defaults to "ros".
Raises:
RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first.
"""
# resolve env_ids
if env_ids is None:
env_ids = self._ALL_INDICES
# get current positions
pos_w, quat_w = self._compute_view_world_poses(env_ids)
if positions is not None:
# transform to camera frame
pos_offset_world_frame = positions - pos_w
self._offset_pos[env_ids] = math_utils.quat_apply(math_utils.quat_inv(quat_w), pos_offset_world_frame)
if orientations is not None:
# convert rotation matrix from input convention to world
quat_w_set = convert_orientation_convention(orientations, origin=convention, target="world")
self._offset_quat[env_ids] = math_utils.quat_mul(math_utils.quat_inv(quat_w), quat_w_set)
# update the data
pos_w, quat_w = self._compute_camera_world_poses(env_ids)
self._data.pos_w[env_ids] = pos_w
self._data.quat_w_world[env_ids] = quat_w
[docs] def set_world_poses_from_view(
self, eyes: torch.Tensor, targets: torch.Tensor, env_ids: Sequence[int] | None = None
):
"""Set the poses of the camera from the eye position and look-at target position.
Args:
eyes: The positions of the camera's eye. Shape is N, 3).
targets: The target locations to look at. Shape is (N, 3).
env_ids: A sensor ids to manipulate. Defaults to None, which means all sensor indices.
Raises:
RuntimeError: If the camera prim is not set. Need to call :meth:`initialize` method first.
NotImplementedError: If the stage up-axis is not "Y" or "Z".
"""
# camera position and rotation in opengl convention
orientations = math_utils.quat_from_matrix(create_rotation_matrix_from_view(eyes, targets, device=self._device))
self.set_world_poses(eyes, orientations, env_ids, convention="opengl")
"""
Implementation.
"""
def _initialize_rays_impl(self):
# Create all indices buffer
self._ALL_INDICES = torch.arange(self._view.count, device=self._device, dtype=torch.long)
# Create frame count buffer
self._frame = torch.zeros(self._view.count, device=self._device, dtype=torch.long)
# create buffers
self._create_buffers()
# compute intrinsic matrices
self._compute_intrinsic_matrices()
# compute ray stars and directions
self.ray_starts, self.ray_directions = self.cfg.pattern_cfg.func(
self.cfg.pattern_cfg, self._data.intrinsic_matrices, self._device
)
self.num_rays = self.ray_directions.shape[1]
# create buffer to store ray hits
self.ray_hits_w = torch.zeros(self._view.count, self.num_rays, 3, device=self._device)
# set offsets
quat_w = convert_orientation_convention(
torch.tensor([self.cfg.offset.rot], device=self._device), origin=self.cfg.offset.convention, target="world"
)
self._offset_quat = quat_w.repeat(self._view.count, 1)
self._offset_pos = torch.tensor(list(self.cfg.offset.pos), device=self._device).repeat(self._view.count, 1)
def _update_buffers_impl(self, env_ids: Sequence[int]):
"""Fills the buffers of the sensor data."""
# increment frame count
self._frame[env_ids] += 1
# compute poses from current view
pos_w, quat_w = self._compute_camera_world_poses(env_ids)
# update the data
self._data.pos_w[env_ids] = pos_w
self._data.quat_w_world[env_ids] = quat_w
# note: full orientation is considered
ray_starts_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_starts[env_ids])
ray_starts_w += pos_w.unsqueeze(1)
ray_directions_w = math_utils.quat_apply(quat_w.repeat(1, self.num_rays), self.ray_directions[env_ids])
# ray cast and store the hits
# TODO: Make ray-casting work for multiple meshes?
# necessary for regular dictionaries.
self.ray_hits_w, ray_depth, ray_normal, _ = raycast_mesh(
ray_starts_w,
ray_directions_w,
mesh=RayCasterCamera.meshes[self.cfg.mesh_prim_paths[0]],
max_dist=self.cfg.max_distance,
return_distance=any(
[name in self.cfg.data_types for name in ["distance_to_image_plane", "distance_to_camera"]]
),
return_normal="normals" in self.cfg.data_types,
)
# update output buffers
if "distance_to_image_plane" in self.cfg.data_types:
# note: data is in camera frame so we only take the first component (z-axis of camera frame)
distance_to_image_plane = (
math_utils.quat_apply(
math_utils.quat_inv(quat_w).repeat(1, self.num_rays),
(ray_depth[:, :, None] * ray_directions_w),
)
)[:, :, 0]
self._data.output["distance_to_image_plane"][env_ids] = distance_to_image_plane.view(-1, *self.image_shape)
if "distance_to_camera" in self.cfg.data_types:
self._data.output["distance_to_camera"][env_ids] = ray_depth.view(-1, *self.image_shape)
if "normals" in self.cfg.data_types:
self._data.output["normals"][env_ids] = ray_normal.view(-1, *self.image_shape, 3)
def _debug_vis_callback(self, event):
# in case it crashes be safe
if not hasattr(self, "ray_hits_w"):
return
# show ray hit positions
self.ray_visualizer.visualize(self.ray_hits_w.view(-1, 3))
"""
Private Helpers
"""
def _check_supported_data_types(self, cfg: RayCasterCameraCfg):
"""Checks if the data types are supported by the ray-caster camera."""
# check if there is any intersection in unsupported types
# reason: we cannot obtain this data from simplified warp-based ray caster
common_elements = set(cfg.data_types) & RayCasterCamera.UNSUPPORTED_TYPES
if common_elements:
raise ValueError(
f"RayCasterCamera class does not support the following sensor types: {common_elements}."
"\n\tThis is because these sensor types cannot be obtained in a fast way using ''warp''."
"\n\tHint: If you need to work with these sensor types, we recommend using the USD camera"
" interface from the omni.isaac.orbit.sensors.camera module."
)
def _create_buffers(self):
"""Create buffers for storing data."""
# prepare drift
self.drift = torch.zeros(self._view.count, 3, device=self.device)
# create the data object
# -- pose of the cameras
self._data.pos_w = torch.zeros((self._view.count, 3), device=self._device)
self._data.quat_w_world = torch.zeros((self._view.count, 4), device=self._device)
# -- intrinsic matrix
self._data.intrinsic_matrices = torch.zeros((self._view.count, 3, 3), device=self._device)
self._data.intrinsic_matrices[:, 2, 2] = 1.0
self._data.image_shape = self.image_shape
# -- output data
# create the buffers to store the annotator data.
self._data.output = TensorDict({}, batch_size=self._view.count, device=self.device)
self._data.info = [{name: None for name in self.cfg.data_types}] * self._view.count
for name in self.cfg.data_types:
if name in ["distance_to_image_plane", "distance_to_camera"]:
shape = (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width)
elif name in ["normals"]:
shape = (self.cfg.pattern_cfg.height, self.cfg.pattern_cfg.width, 3)
else:
raise ValueError(f"Received unknown data type: {name}. Please check the configuration.")
# allocate tensor to store the data
self._data.output[name] = torch.zeros((self._view.count, *shape), device=self._device)
def _compute_intrinsic_matrices(self):
"""Computes the intrinsic matrices for the camera based on the config provided."""
# get the sensor properties
pattern_cfg = self.cfg.pattern_cfg
# compute the intrinsic matrix
vertical_aperture = pattern_cfg.horizontal_aperture * pattern_cfg.height / pattern_cfg.width
f_x = pattern_cfg.width * pattern_cfg.focal_length / pattern_cfg.horizontal_aperture
f_y = pattern_cfg.height * pattern_cfg.focal_length / vertical_aperture
c_x = pattern_cfg.horizontal_aperture_offset * f_x + pattern_cfg.width / 2
c_y = pattern_cfg.vertical_aperture_offset * f_y + pattern_cfg.height / 2
# allocate the intrinsic matrices
self._data.intrinsic_matrices[:, 0, 0] = f_x
self._data.intrinsic_matrices[:, 0, 2] = c_x
self._data.intrinsic_matrices[:, 1, 1] = f_y
self._data.intrinsic_matrices[:, 1, 2] = c_y
# save focal length
self._focal_length = pattern_cfg.focal_length
def _compute_view_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]:
"""Obtains the pose of the view the camera is attached to in the world frame.
Returns:
A tuple of the position (in meters) and quaternion (w, x, y, z).
"""
# obtain the poses of the sensors
# note: clone arg doesn't exist for xform prim view so we need to do this manually
if isinstance(self._view, XFormPrimView):
pos_w, quat_w = self._view.get_world_poses(env_ids)
elif isinstance(self._view, physx.ArticulationView):
pos_w, quat_w = self._view.get_root_transforms()[env_ids].split([3, 4], dim=-1)
quat_w = math_utils.convert_quat(quat_w, to="wxyz")
elif isinstance(self._view, physx.RigidBodyView):
pos_w, quat_w = self._view.get_transforms()[env_ids].split([3, 4], dim=-1)
quat_w = math_utils.convert_quat(quat_w, to="wxyz")
else:
raise RuntimeError(f"Unsupported view type: {type(self._view)}")
# return the pose
return pos_w.clone(), quat_w.clone()
def _compute_camera_world_poses(self, env_ids: Sequence[int]) -> tuple[torch.Tensor, torch.Tensor]:
"""Computes the pose of the camera in the world frame.
This function applies the offset pose to the pose of the view the camera is attached to.
Returns:
A tuple of the position (in meters) and quaternion (w, x, y, z) in "world" convention.
"""
# get the pose of the view the camera is attached to
pos_w, quat_w = self._compute_view_world_poses(env_ids)
# apply offsets
# need to apply quat because offset relative to parent frame
pos_w += math_utils.quat_apply(quat_w, self._offset_pos[env_ids])
quat_w = math_utils.quat_mul(quat_w, self._offset_quat[env_ids])
return pos_w, quat_w
