hofee/gen_data_curobo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

constrained planning, robot segmentation

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2024-02-22 21:45:47 -08:00
parent 88eac64edc
commit bafdf80c05
102 changed files with 12440 additions and 8112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
src/curobo/geom/cv.py Normal file
View File

@@ -0,0 +1,118 @@
#
# 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
@torch.jit.script
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
@torch.jit.script
def get_projection_rays(height: int, width: int, 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]
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 * (1.0 / 1000.0)
return rays
@torch.jit.script
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
@torch.jit.script
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

23
src/curobo/geom/sdf/world.py
View File

@@ -364,8 +364,12 @@ class WorldPrimitiveCollision(WorldCollision):
if self.cache is not None and "obb" in self.cache and self.cache["obb"] not in [None, 0]:
self._create_obb_cache(self.cache["obb"])
def load_collision_model(self, world_config: WorldConfig, env_idx=0):
self._load_collision_model_in_cache(world_config, env_idx)
def load_collision_model(
self, world_config: WorldConfig, env_idx=0, fix_cache_reference: bool = False
):
self._load_collision_model_in_cache(
world_config, env_idx, fix_cache_reference=fix_cache_reference
)
def load_batch_collision_model(self, world_config_list: List[WorldConfig]):
"""Load a batch of collision environments from a list of world configs.
@@ -436,7 +440,9 @@ class WorldPrimitiveCollision(WorldCollision):
)
self.collision_types["primitive"] = True
def _load_collision_model_in_cache(self, world_config: WorldConfig, env_idx: int = 0):
def _load_collision_model_in_cache(
self, world_config: WorldConfig, env_idx: int = 0, fix_cache_reference: bool = False
):
cube_objs = world_config.cuboid
max_obb = len(cube_objs)
self.world_model = world_config
@@ -444,8 +450,11 @@ class WorldPrimitiveCollision(WorldCollision):
log_info("No OBB objs")
return
if self._cube_tensor_list is None or self._cube_tensor_list[0].shape[1] < max_obb:
log_info("Creating Obb cache" + str(max_obb))
self._create_obb_cache(max_obb)
if not fix_cache_reference:
log_info("Creating Obb cache" + str(max_obb))
self._create_obb_cache(max_obb)
else:
log_error("number of OBB is larger than collision cache, create larger cache.")
# load as a batch:
pose_batch = [c.pose for c in cube_objs]
@@ -835,7 +844,9 @@ class WorldPrimitiveCollision(WorldCollision):
return dist
def clear_cache(self):
pass
if self._cube_tensor_list is not None:
self._cube_tensor_list[2][:] = 1
self._env_n_obbs[:] = 0
def get_voxels_in_bounding_box(
self,

11
src/curobo/geom/sdf/world_blox.py
View File

@@ -55,7 +55,7 @@ class WorldBloxCollision(WorldMeshCollision):
self._blox_voxel_sizes = [0.02]
super().__init__(config)
def load_collision_model(self, world_model: WorldConfig):
def load_collision_model(self, world_model: WorldConfig, fix_cache_reference: bool = False):
# load nvblox mesh
if len(world_model.blox) > 0:
# check if there is a mapper instance:
@@ -109,15 +109,17 @@ class WorldBloxCollision(WorldMeshCollision):
self._blox_names = names
self.collision_types["blox"] = True
return super().load_collision_model(world_model)
return super().load_collision_model(world_model, fix_cache_reference=fix_cache_reference)
def clear_cache(self):
self._blox_mapper.clear()
self._blox_mapper.update_hashmaps()
super().clear_cache()
def clear_blox_layer(self, layer_name: str):
index = self._blox_names.index(layer_name)
self._blox_mapper.clear(index)
self._blox_mapper.update_hashmaps()
def _get_blox_sdf(
self,
@@ -147,6 +149,7 @@ class WorldBloxCollision(WorldMeshCollision):
speed_dt,
sweep_steps,
enable_speed_metric,
return_loss=False,
):
d = self._blox_mapper.query_sphere_trajectory_sdf_cost(
query_spheres,
@@ -160,6 +163,8 @@ class WorldBloxCollision(WorldMeshCollision):
self._blox_tensor_list[1],
sweep_steps,
enable_speed_metric,
return_loss,
use_experimental=False,
)
return d
@@ -279,6 +284,7 @@ class WorldBloxCollision(WorldMeshCollision):
speed_dt=speed_dt,
sweep_steps=sweep_steps,
enable_speed_metric=enable_speed_metric,
return_loss=return_loss,
)
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (
@@ -332,6 +338,7 @@ class WorldBloxCollision(WorldMeshCollision):
speed_dt=speed_dt,
sweep_steps=sweep_steps,
enable_speed_metric=enable_speed_metric,
return_loss=return_loss,
)
if ("primitive" not in self.collision_types or not self.collision_types["primitive"]) and (

12
src/curobo/geom/sdf/world_mesh.py
View File

@@ -72,7 +72,11 @@ class WorldMeshCollision(WorldPrimitiveCollision):
return super()._init_cache()
def load_collision_model(
self, world_model: WorldConfig, env_idx: int = 0, load_obb_obs: bool = True
self,
world_model: WorldConfig,
env_idx: int = 0,
load_obb_obs: bool = True,
fix_cache_reference: bool = False,
):
max_nmesh = len(world_model.mesh)
if max_nmesh > 0:
@@ -91,14 +95,16 @@ class WorldMeshCollision(WorldPrimitiveCollision):
self.collision_types["mesh"] = True
if load_obb_obs:
super().load_collision_model(world_model, env_idx)
super().load_collision_model(
world_model, env_idx, fix_cache_reference=fix_cache_reference
)
else:
self.world_model = world_model
def load_batch_collision_model(self, world_config_list: List[WorldConfig]):
max_nmesh = max([len(x.mesh) for x in world_config_list])
if self._mesh_tensor_list is None or self._mesh_tensor_list[0].shape[1] < max_nmesh:
log_info("Creating new Mesh cache: " + str(max_nmesh))
log_warn("Creating new Mesh cache: " + str(max_nmesh))
self._create_mesh_cache(max_nmesh)
for env_idx, world_model in enumerate(world_config_list):

113
src/curobo/geom/transform.py
View File

@@ -228,6 +228,61 @@ def pose_inverse(
return out_position, out_quaternion
@wp.kernel
def compute_pose_inverse(
position: wp.array(dtype=wp.vec3),
quat: wp.array(dtype=wp.vec4),
out_position: wp.array(dtype=wp.vec3),
out_quat: wp.array(dtype=wp.vec4),
): # b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_position = position[b_idx]
in_quat = quat[b_idx]
# read point
# create a transform from a vector/quaternion:
t_1 = wp.transform(in_position, wp.quaternion(in_quat[1], in_quat[2], in_quat[3], in_quat[0]))
t_3 = wp.transform_inverse(t_1)
# write pt:
out_q = wp.transform_get_rotation(t_3)
out_v = wp.vec4()
out_v[0] = out_q[3] # out_q[3]
out_v[1] = out_q[0] # [0]
out_v[2] = out_q[1] # wp.extract(out_q, 1)
out_v[3] = out_q[2] # wp.extract(out_q, 2)
out_position[b_idx] = wp.transform_get_translation(t_3)
out_quat[b_idx] = out_v
@wp.kernel
def compute_matrix_to_quat(
in_mat: wp.array(dtype=wp.mat33),
out_quat: wp.array(dtype=wp.vec4),
):
# b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_m = in_mat[b_idx]
# read point
# create a transform from a vector/quaternion:
out_q = wp.quat_from_matrix(in_m)
out_v = wp.vec4()
out_v[0] = out_q[3] # wp.extract(out_q, 3)
out_v[1] = out_q[0] # wp.extract(out_q, 0)
out_v[2] = out_q[1] # wp.extract(out_q, 1)
out_v[3] = out_q[2] # wp.extract(out_q, 2)
# write pt:
out_quat[b_idx] = out_v
@wp.kernel
def compute_transform_point(
position: wp.array(dtype=wp.vec3),
@@ -331,37 +386,6 @@ def compute_batch_pose_multipy(
out_quat[b_idx] = out_v
@wp.kernel
def compute_pose_inverse(
position: wp.array(dtype=wp.vec3),
quat: wp.array(dtype=wp.vec4),
out_position: wp.array(dtype=wp.vec3),
out_quat: wp.array(dtype=wp.vec4),
): # b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_position = position[b_idx]
in_quat = quat[b_idx]
# read point
# create a transform from a vector/quaternion:
t_1 = wp.transform(in_position, wp.quaternion(in_quat[1], in_quat[2], in_quat[3], in_quat[0]))
t_3 = wp.transform_inverse(t_1)
# write pt:
out_q = wp.transform_get_rotation(t_3)
out_v = wp.vec4()
out_v[0] = out_q[3]
out_v[1] = out_q[0]
out_v[2] = out_q[1]
out_v[3] = out_q[2]
out_position[b_idx] = wp.transform_get_translation(t_3)
out_quat[b_idx] = out_v
@wp.kernel
def compute_quat_to_matrix(
quat: wp.array(dtype=wp.vec4),
@@ -382,30 +406,6 @@ def compute_quat_to_matrix(
out_mat[b_idx] = m_1
@wp.kernel
def compute_matrix_to_quat(
in_mat: wp.array(dtype=wp.mat33),
out_quat: wp.array(dtype=wp.vec4),
):
# b pose_1 and b pose_2, compute pose_1 * pose_2
b_idx = wp.tid()
# read data:
in_m = in_mat[b_idx]
# read point
# create a transform from a vector/quaternion:
out_q = wp.quat_from_matrix(in_m)
out_v = wp.vec4()
out_v[0] = out_q[3]
out_v[1] = out_q[0]
out_v[2] = out_q[1]
out_v[3] = out_q[2]
# write pt:
out_quat[b_idx] = out_v
@wp.kernel
def compute_pose_multipy(
position: wp.array(dtype=wp.vec3),
@@ -962,6 +962,7 @@ class PoseInverse(torch.autograd.Function):
adj_quaternion,
)
ctx.b = b
wp.launch(
kernel=compute_pose_inverse,
dim=b,
@@ -1071,6 +1072,7 @@ class QuatToMatrix(torch.autograd.Function):
adj_quaternion,
)
ctx.b = b
wp.launch(
kernel=compute_quat_to_matrix,
dim=b,
@@ -1153,6 +1155,7 @@ class MatrixToQuaternion(torch.autograd.Function):
adj_mat,
)
ctx.b = b
wp.launch(
kernel=compute_matrix_to_quat,
dim=b,

135
src/curobo/geom/types.py
View File

@@ -12,7 +12,7 @@ from __future__ import annotations
# Standard Library
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence
from typing import Any, Dict, List, Optional, Sequence, Union
# Third Party
import numpy as np
@@ -22,6 +22,7 @@ import trimesh
# CuRobo
from curobo.geom.sphere_fit import SphereFitType, fit_spheres_to_mesh
from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.util.logger import log_error, log_warn
from curobo.util_file import get_assets_path, join_path
@@ -60,7 +61,7 @@ class Obstacle:
tensor_args: TensorDeviceType = TensorDeviceType()
def get_trimesh_mesh(self, process: bool = True) -> trimesh.Trimesh:
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True) -> trimesh.Trimesh:
"""Create a trimesh instance from the obstacle representation.
Args:
@@ -74,8 +75,11 @@ class Obstacle:
"""
raise NotImplementedError
def save_as_mesh(self, file_path: str):
def save_as_mesh(self, file_path: str, transform_with_pose: bool = False):
mesh_scene = self.get_trimesh_mesh()
if transform_with_pose:
mesh_scene.apply_transform(self.get_transform_matrix())
mesh_scene.export(file_path)
def get_cuboid(self) -> Cuboid:
@@ -238,7 +242,7 @@ class Cuboid(Obstacle):
if self.pose is None:
log_error("Cuboid Obstacle requires Pose")
def get_trimesh_mesh(self, process: bool = True):
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
m = trimesh.creation.box(extents=self.dims)
if self.color is not None:
color_visual = trimesh.visual.color.ColorVisuals(
@@ -254,7 +258,7 @@ class Capsule(Obstacle):
base: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
tip: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
def get_trimesh_mesh(self, process: bool = True):
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
height = self.tip[2] - self.base[2]
if (
height < 0
@@ -280,7 +284,7 @@ class Cylinder(Obstacle):
radius: float = 0.0
height: float = 0.0
def get_trimesh_mesh(self, process: bool = True):
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
m = trimesh.creation.cylinder(radius=self.radius, height=self.height)
if self.color is not None:
color_visual = trimesh.visual.color.ColorVisuals(
@@ -304,7 +308,7 @@ class Sphere(Obstacle):
if self.pose is not None:
self.position = self.pose[:3]
def get_trimesh_mesh(self, process: bool = True):
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
m = trimesh.creation.icosphere(radius=self.radius)
if self.color is not None:
color_visual = trimesh.visual.color.ColorVisuals(
@@ -356,9 +360,9 @@ class Sphere(Obstacle):
@dataclass
class Mesh(Obstacle):
file_path: Optional[str] = None
file_string: Optional[
str
] = None # storing full mesh as a string, loading from this is not implemented yet.
file_string: Optional[str] = (
None # storing full mesh as a string, loading from this is not implemented yet.
)
urdf_path: Optional[str] = None # useful for visualization in isaac gym.
vertices: Optional[List[List[float]]] = None
faces: Optional[List[int]] = None
@@ -375,13 +379,13 @@ class Mesh(Obstacle):
self.vertices = np.ravel(self.scale) * self.vertices
self.scale = None
def get_trimesh_mesh(self, process: bool = True):
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
# load mesh from filepath or verts and faces:
if self.file_path is not None:
m = trimesh.load(self.file_path, process=process, force="mesh")
if isinstance(m, trimesh.Scene):
m = m.dump(concatenate=True)
if isinstance(m.visual, trimesh.visual.texture.TextureVisuals):
if process_color and isinstance(m.visual, trimesh.visual.texture.TextureVisuals):
m.visual = m.visual.to_color()
else:
m = trimesh.Trimesh(
@@ -467,7 +471,7 @@ class BloxMap(Obstacle):
name=self.name + "_mesh", file_path=self.mesh_file_path, pose=self.pose
)
def get_trimesh_mesh(self, process: bool = True):
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
if self.mesh is not None:
return self.mesh.get_trimesh_mesh(process)
else:
@@ -475,6 +479,91 @@ class BloxMap(Obstacle):
return None
@dataclass
class PointCloud(Obstacle):
points: Union[torch.Tensor, np.ndarray, List[List[float]]] = None
points_features: Union[torch.Tensor, np.ndarray, List[List[float]], None] = None
def __post_init__(self):
if self.scale is not None and self.points is not None:
self.points = np.ravel(self.scale) * self.points
self.scale = None
def get_trimesh_mesh(self, process: bool = True, process_color: bool = True):
points = self.points
if isinstance(points, torch.Tensor):
points = points.view(-1, 3).cpu().numpy()
if isinstance(points, list):
points = np.ndarray(points)
mesh = Mesh.from_pointcloud(points, pose=self.pose)
return mesh.get_trimesh_mesh()
def get_mesh_data(self, process: bool = True):
verts = faces = None
m = self.get_trimesh_mesh(process=process)
verts = m.vertices.view(np.ndarray)
faces = m.faces
return verts, faces
@staticmethod
def from_camera_observation(
camera_obs: CameraObservation, name: str = "pc_obstacle", pose: Optional[List[float]] = None
):
return PointCloud(name=name, pose=pose, points=camera_obs.get_pointcloud())
def get_bounding_spheres(
self,
n_spheres: int = 1,
surface_sphere_radius: float = 0.002,
fit_type: SphereFitType = SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
voxelize_method: str = "ray",
pre_transform_pose: Optional[Pose] = None,
tensor_args: TensorDeviceType = TensorDeviceType(),
) -> List[Sphere]:
"""Compute n spheres that fits in the volume of the object.
Args:
n: number of spheres
Returns:
spheres
"""
# sample points in pointcloud:
# mesh = self.get_trimesh_mesh()
# pts, n_radius = fit_spheres_to_mesh(
# mesh, n_spheres, surface_sphere_radius, fit_type, voxelize_method=voxelize_method
# )
obj_pose = Pose.from_list(self.pose, tensor_args)
# transform object:
# transform points:
if pre_transform_pose is not None:
obj_pose = pre_transform_pose.multiply(obj_pose) # convert object pose to another frame
if pts is None or len(pts) == 0:
log_warn("spheres could not be fit!, adding one sphere at origin")
pts = np.zeros((1, 3))
pts[0, :] = mesh.centroid
n_radius = [0.02]
obj_pose = Pose.from_list([0, 0, 0, 1, 0, 0, 0], tensor_args)
points_cuda = tensor_args.to_device(pts)
pts = obj_pose.transform_points(points_cuda).cpu().view(-1, 3).numpy()
new_spheres = [
Sphere(
name="sph_" + str(i),
pose=[pts[i, 0], pts[i, 1], pts[i, 2], 1, 0, 0, 0],
radius=n_radius[i],
)
for i in range(pts.shape[0])
]
return new_spheres
@dataclass
class WorldConfig(Sequence):
"""Representation of World for use in CuRobo."""
@@ -664,23 +753,25 @@ class WorldConfig(Sequence):
)
@staticmethod
def get_scene_graph(current_world: WorldConfig):
def get_scene_graph(current_world: WorldConfig, process_color: bool = True):
m_world = WorldConfig.create_mesh_world(current_world)
mesh_scene = trimesh.scene.scene.Scene(base_frame="world")
mesh_scene = trimesh.scene.scene.Scene(base_frame="world_origin")
mesh_list = m_world
for m in mesh_list:
mesh_scene.add_geometry(
m.get_trimesh_mesh(),
m.get_trimesh_mesh(process_color=process_color),
geom_name=m.name,
parent_node_name="world",
parent_node_name="world_origin",
transform=m.get_transform_matrix(),
)
return mesh_scene
@staticmethod
def create_merged_mesh_world(current_world: WorldConfig, process: bool = True):
mesh_scene = WorldConfig.get_scene_graph(current_world)
def create_merged_mesh_world(
current_world: WorldConfig, process: bool = True, process_color: bool = True
):
mesh_scene = WorldConfig.get_scene_graph(current_world, process_color=process_color)
mesh_scene = mesh_scene.dump(concatenate=True)
new_mesh = Mesh(
vertices=mesh_scene.vertices.view(np.ndarray),
@@ -702,8 +793,10 @@ class WorldConfig(Sequence):
def get_collision_check_world(self, mesh_process: bool = False):
return WorldConfig.create_collision_support_world(self, process=mesh_process)
def save_world_as_mesh(self, file_path: str, save_as_scene_graph=False):
mesh_scene = WorldConfig.get_scene_graph(self)
def save_world_as_mesh(
self, file_path: str, save_as_scene_graph=False, process_color: bool = True
):
mesh_scene = WorldConfig.get_scene_graph(self, process_color=process_color)
if save_as_scene_graph:
mesh_scene = mesh_scene.dump(concatenate=True)