Significantly improved convergence for mesh and cuboid, new ESDF collision.

This commit is contained in:
Balakumar Sundaralingam
2024-03-18 11:19:48 -07:00
parent 286b3820a5
commit b1f63e8778
100 changed files with 7587 additions and 2589 deletions

View File

@@ -12,8 +12,11 @@
# Third Party
import torch
# CuRobo
from curobo.util.torch_utils import get_torch_jit_decorator
@torch.jit.script
@get_torch_jit_decorator()
def project_depth_to_pointcloud(depth_image: torch.Tensor, intrinsics_matrix: torch.Tensor):
"""Projects numpy depth image to point cloud.
@@ -43,7 +46,7 @@ def project_depth_to_pointcloud(depth_image: torch.Tensor, intrinsics_matrix: to
return raw_pc
@torch.jit.script
@get_torch_jit_decorator()
def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor):
"""Projects numpy depth image to point cloud.
@@ -54,10 +57,10 @@ def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor
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]
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)
@@ -73,7 +76,6 @@ def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor
device=intrinsics_matrix.device,
dtype=torch.float32,
)
output_x = (input_x - cx) / fx
output_y = (input_y - cy) / fy
@@ -84,7 +86,7 @@ def get_projection_rays(height: int, width: int, intrinsics_matrix: torch.Tensor
return rays
@torch.jit.script
@get_torch_jit_decorator()
def project_pointcloud_to_depth(
pointcloud: torch.Tensor,
output_image: torch.Tensor,
@@ -106,7 +108,7 @@ def project_pointcloud_to_depth(
return output_image
@torch.jit.script
@get_torch_jit_decorator()
def project_depth_using_rays(
depth_image: torch.Tensor, rays: torch.Tensor, filter_origin: bool = False
):