123 lines
3.8 KiB
Python
123 lines
3.8 KiB
Python
#
|
|
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
#
|
|
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
|
# property and proprietary rights in and to this material, related
|
|
# documentation and any modifications thereto. Any use, reproduction,
|
|
# disclosure or distribution of this material and related documentation
|
|
# without an express license agreement from NVIDIA CORPORATION or
|
|
# its affiliates is strictly prohibited.
|
|
#
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# CuRobo
|
|
from curobo.util.torch_utils import get_torch_jit_decorator
|
|
|
|
|
|
@get_torch_jit_decorator()
|
|
def project_depth_to_pointcloud(depth_image: torch.Tensor, intrinsics_matrix: torch.Tensor):
|
|
"""Projects numpy depth image to point cloud.
|
|
|
|
Args:
|
|
np_depth_image: numpy array float, shape (h, w).
|
|
intrinsics array: numpy array float, 3x3 intrinsics matrix.
|
|
|
|
Returns:
|
|
array of float (h, w, 3)
|
|
"""
|
|
fx = intrinsics_matrix[0, 0]
|
|
fy = intrinsics_matrix[1, 1]
|
|
cx = intrinsics_matrix[0, 2]
|
|
cy = intrinsics_matrix[1, 2]
|
|
height = depth_image.shape[0]
|
|
width = depth_image.shape[1]
|
|
input_x = torch.arange(width, dtype=torch.float32, device=depth_image.device)
|
|
input_y = torch.arange(height, dtype=torch.float32, device=depth_image.device)
|
|
input_x, input_y = torch.meshgrid(input_x, input_y, indexing="ij")
|
|
input_x, input_y = input_x.T, input_y.T
|
|
|
|
input_z = depth_image
|
|
output_x = (input_x * input_z - cx * input_z) / fx
|
|
output_y = (input_y * input_z - cy * input_z) / fy
|
|
raw_pc = torch.stack([output_x, output_y, input_z], -1)
|
|
|
|
return raw_pc
|
|
|
|
|
|
@get_torch_jit_decorator()
|
|
def get_projection_rays(
|
|
height: int, width: int, intrinsics_matrix: torch.Tensor, depth_to_meter: float = 0.001
|
|
):
|
|
"""Projects numpy depth image to point cloud.
|
|
|
|
Args:
|
|
np_depth_image: numpy array float, shape (h, w).
|
|
intrinsics array: numpy array float, 3x3 intrinsics matrix.
|
|
|
|
Returns:
|
|
array of float (h, w, 3)
|
|
"""
|
|
fx = intrinsics_matrix[:, 0:1, 0:1]
|
|
fy = intrinsics_matrix[:, 1:2, 1:2]
|
|
cx = intrinsics_matrix[:, 0:1, 2:3]
|
|
cy = intrinsics_matrix[:, 1:2, 2:3]
|
|
|
|
input_x = torch.arange(width, dtype=torch.float32, device=intrinsics_matrix.device)
|
|
input_y = torch.arange(height, dtype=torch.float32, device=intrinsics_matrix.device)
|
|
input_x, input_y = torch.meshgrid(input_x, input_y, indexing="ij")
|
|
|
|
input_x, input_y = input_x.T, input_y.T
|
|
|
|
input_x = input_x.unsqueeze(0).repeat(intrinsics_matrix.shape[0], 1, 1)
|
|
input_y = input_y.unsqueeze(0).repeat(intrinsics_matrix.shape[0], 1, 1)
|
|
|
|
input_z = torch.ones(
|
|
(intrinsics_matrix.shape[0], height, width),
|
|
device=intrinsics_matrix.device,
|
|
dtype=torch.float32,
|
|
)
|
|
output_x = (input_x - cx) / fx
|
|
output_y = (input_y - cy) / fy
|
|
|
|
rays = torch.stack([output_x, output_y, input_z], -1).reshape(
|
|
intrinsics_matrix.shape[0], width * height, 3
|
|
)
|
|
rays = rays * depth_to_meter
|
|
return rays
|
|
|
|
|
|
@get_torch_jit_decorator()
|
|
def project_pointcloud_to_depth(
|
|
pointcloud: torch.Tensor,
|
|
output_image: torch.Tensor,
|
|
):
|
|
"""Projects pointcloud to depth image
|
|
|
|
Args:
|
|
np_depth_image: numpy array float, shape (h, w).
|
|
intrinsics array: numpy array float, 3x3 intrinsics matrix.
|
|
|
|
Returns:
|
|
array of float (h, w)
|
|
"""
|
|
width, height = output_image.shape
|
|
|
|
output_image = output_image.view(-1)
|
|
output_image[:] = pointcloud[:, 2]
|
|
output_image = output_image.view(width, height)
|
|
return output_image
|
|
|
|
|
|
@get_torch_jit_decorator()
|
|
def project_depth_using_rays(
|
|
depth_image: torch.Tensor, rays: torch.Tensor, filter_origin: bool = False
|
|
):
|
|
if filter_origin:
|
|
depth_image = torch.where(depth_image < 0.01, 0, depth_image)
|
|
|
|
depth_image = depth_image.view(depth_image.shape[0], -1, 1).contiguous()
|
|
points = depth_image * rays
|
|
return points
|