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

Bug Fixes, Add Isaac Sim 4.5 support for examples

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2025-04-25 11:24:16 -07:00
parent 2fbffc3522
commit 0a50de1ba7
43 changed files with 728 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
examples/collision_check_example.py
View File

@@ -8,9 +8,7 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
""" Example computing collisions using curobo
"""
"""Example computing collisions using curobo"""
# Third Party
import torch

8
examples/isaac_sim/batch_collision_checker.py
View File

@@ -51,7 +51,13 @@ import carb
import numpy as np
from helper import add_extensions
from omni.isaac.core import World
from omni.isaac.core.materials import OmniPBR
try:
from omni.isaac.core.materials import OmniPBR
except ImportError:
from isaacsim.core.api.materials import OmniPBR
from omni.isaac.core.objects import sphere
# CuRobo

9
examples/isaac_sim/batch_motion_gen_reacher.py
View File

@@ -126,9 +126,11 @@ def main():
"/World/world_" + str(i) + "/",
robot_name="robot_" + str(i),
position=pose.position[0].cpu().numpy(),
initialize_world=False,
)
robot_list.append(r[0])
setup_curobo_logger("warn")
my_world.initialize_physics()
# warmup curobo instance
@@ -195,9 +197,10 @@ def main():
continue
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
for robot in robot_list:
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -223,6 +226,8 @@ def main():
past_goal = ik_goal.clone()
sim_js_names = robot_list[0].dof_names
sim_js = robot_list[0].get_joints_state()
if sim_js is None:
continue
full_js = JointState(
position=tensor_args.to_device(sim_js.positions).view(1, -1),
velocity=tensor_args.to_device(sim_js.velocities).view(1, -1) * 0.0,

16
examples/isaac_sim/collision_checker.py
View File

@@ -54,7 +54,11 @@ import carb
import numpy as np
from helper import add_extensions
from omni.isaac.core import World
from omni.isaac.core.materials import OmniPBR
try:
from omni.isaac.core.materials import OmniPBR
except ImportError:
from isaacsim.core.api.materials import OmniPBR
from omni.isaac.core.objects import sphere
# CuRobo
@@ -73,7 +77,10 @@ from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
def draw_line(start, gradient):
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:
@@ -194,7 +201,10 @@ def main():
else:
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:

13
examples/isaac_sim/constrained_reacher.py
View File

@@ -58,7 +58,11 @@ import argparse
# Third Party
import carb
from omni.isaac.core import World
from omni.isaac.core.materials import OmniGlass, OmniPBR
try:
from omni.isaac.core.materials import OmniGlass, OmniPBR
except ImportError:
from isaacsim.core.api.materials import OmniGlass, OmniPBR
from omni.isaac.core.objects import cuboid, sphere
from omni.isaac.core.utils.types import ArticulationAction
@@ -207,8 +211,9 @@ if __name__ == "__main__":
continue
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -218,7 +223,7 @@ if __name__ == "__main__":
if False and step_index % 50 == 0.0: # No obstacle update
obstacles = usd_help.get_obstacles_from_stage(
# only_paths=[obstacles_path],
only_paths=["/World"],
reference_prim_path=robot_prim_path,
ignore_substring=[
robot_prim_path,

80
examples/isaac_sim/helper.py
View File

@@ -16,7 +16,6 @@ from typing import Dict, List
import numpy as np
from matplotlib import cm
from omni.isaac.core import World
from omni.isaac.core.materials import OmniPBR
from omni.isaac.core.objects import cuboid
from omni.isaac.core.robots import Robot
from pxr import UsdPhysics
@@ -26,14 +25,28 @@ from curobo.util.logger import log_warn
from curobo.util.usd_helper import set_prim_transform
ISAAC_SIM_23 = False
ISAAC_SIM_45 = False
try:
# Third Party
from omni.isaac.urdf import _urdf # isaacsim 2022.2
except ImportError:
# Third Party
from omni.importer.urdf import _urdf # isaac sim 2023.1 or above
try:
from omni.importer.urdf import _urdf # isaac sim 2023.1 or above
except ImportError:
from isaacsim.asset.importer.urdf import _urdf # isaac sim 4.5+
ISAAC_SIM_45 = True
ISAAC_SIM_23 = True
try:
# import for older isaacsim installations
from omni.isaac.core.materials import OmniPBR
except ImportError:
# import for isaac sim 4.5+
from isaacsim.core.api.materials import OmniPBR
# Standard Library
from typing import Optional
@@ -67,6 +80,7 @@ def add_robot_to_scene(
subroot: str = "",
robot_name: str = "robot",
position: np.array = np.array([0, 0, 0]),
initialize_world: bool = True,
):
urdf_interface = _urdf.acquire_urdf_interface()
@@ -91,18 +105,57 @@ def add_robot_to_scene(
and robot_config["kinematics"]["external_asset_path"] is not None
):
asset_path = robot_config["kinematics"]["external_asset_path"]
# urdf_path:
# meshes_path:
# meshes path should be a subset of urdf_path
full_path = join_path(asset_path, robot_config["kinematics"]["urdf_path"])
# full path contains the path to urdf
# Get meshes path
robot_path = get_path_of_dir(full_path)
filename = get_filename(full_path)
imported_robot = urdf_interface.parse_urdf(robot_path, filename, import_config)
dest_path = subroot
robot_path = urdf_interface.import_robot(
robot_path,
filename,
imported_robot,
import_config,
dest_path,
)
if ISAAC_SIM_45:
from isaacsim.core.utils.extensions import get_extension_path_from_name
import omni.kit.commands
import omni.usd
# Retrieve the path of the URDF file from the extension
extension_path = get_extension_path_from_name("isaacsim.asset.importer.urdf")
root_path = robot_path
file_name = filename
# Parse the robot's URDF file to generate a robot model
dest_path = join_path(
root_path, get_filename(file_name, remove_extension=True) + "_temp.usd"
)
result, robot_path = omni.kit.commands.execute(
"URDFParseAndImportFile",
urdf_path="{}/{}".format(root_path, file_name),
import_config=import_config,
dest_path=dest_path,
)
prim_path = omni.usd.get_stage_next_free_path(
my_world.scene.stage,
str(my_world.scene.stage.GetDefaultPrim().GetPath()) + robot_path,
False,
)
robot_prim = my_world.scene.stage.OverridePrim(prim_path)
robot_prim.GetReferences().AddReference(dest_path)
robot_path = prim_path
else:
imported_robot = urdf_interface.parse_urdf(robot_path, filename, import_config)
dest_path = subroot
robot_path = urdf_interface.import_robot(
robot_path,
filename,
imported_robot,
import_config,
dest_path,
)
base_link_name = robot_config["kinematics"]["base_link"]
@@ -117,6 +170,11 @@ def add_robot_to_scene(
set_prim_transform(linkp, [position[0], position[1], position[2], 1, 0, 0, 0])
robot = my_world.scene.add(robot_p)
if initialize_world:
if ISAAC_SIM_45:
my_world.initialize_physics()
robot.initialize()
return robot, robot_path

12
examples/isaac_sim/ik_reachability.py
View File

@@ -112,7 +112,10 @@ def get_pose_grid(n_x, n_y, n_z, max_x, max_y, max_z):
def draw_points(pose, success):
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
N = 100
@@ -236,9 +239,8 @@ def main():
continue
step_index = my_world.current_time_step_index
# print(step_index)
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -251,7 +253,7 @@ def main():
if step_index == 50 or step_index % 500 == 0.0: # and cmd_plan is None:
print("Updating world, reading w.r.t.", robot_prim_path)
obstacles = usd_help.get_obstacles_from_stage(
# only_paths=[obstacles_path],
only_paths=["/World"],
reference_prim_path=robot_prim_path,
ignore_substring=[
robot_prim_path,

15
examples/isaac_sim/load_all_robots.py
View File

@@ -68,13 +68,17 @@ from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
########### OV #################
from curobo.util.logger import setup_curobo_logger
from curobo.util.logger import setup_curobo_logger, log_warn
from curobo.util.usd_helper import UsdHelper
from curobo.util_file import get_motion_gen_robot_list, get_robot_configs_path, join_path, load_yaml
def main():
log_warn("This example will not work correctly in isaac sim 4.5+")
list_of_robots = get_motion_gen_robot_list() # [:2]
# remove tm12 as meshes don't load correctly in isaac sim 4.5
if "tm12.yml" in list_of_robots:
list_of_robots.remove("tm12.yml")
usd_help = UsdHelper()
# assuming obstacles are in objects_path:
@@ -107,12 +111,14 @@ def main():
pose.position[0, 1] = 0
robot_cfg = load_yaml(join_path(get_robot_configs_path(), list_of_robots[i]))["robot_cfg"]
robot_cfg_list.append(robot_cfg)
print("Loading robot: ", list_of_robots[i])
r = add_robot_to_scene(
robot_cfg,
my_world,
"/World/world_" + str(i) + "/",
robot_name="/World/world_" + str(i) + "/" "robot_" + str(i),
position=pose.position[0].cpu().numpy(),
initialize_world=False,
)
robot_list.append(r[0])
@@ -135,6 +141,9 @@ def main():
setup_curobo_logger("warn")
my_world.initialize_physics()
my_world.reset()
add_extensions(simulation_app, args.headless_mode)
while simulation_app.is_running():
my_world.step(render=True)
@@ -146,9 +155,9 @@ def main():
continue
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
if step_index <= 10:
for ri, robot in enumerate(robot_list):
robot._articulation_view.initialize()
j_names = robot_cfg_list[ri]["kinematics"]["cspace"]["joint_names"]
default_config = robot_cfg_list[ri]["kinematics"]["cspace"]["retract_config"]

12
examples/isaac_sim/motion_gen_reacher.py
View File

@@ -272,12 +272,9 @@ def main():
continue
step_index = my_world.current_time_step_index
# print(step_index)
if articulation_controller is None:
# robot.initialize()
articulation_controller = robot.get_articulation_controller()
if step_index < 2:
my_world.reset()
if step_index < 10:
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -291,7 +288,7 @@ def main():
if step_index == 50 or step_index % 1000 == 0.0:
print("Updating world, reading w.r.t.", robot_prim_path)
obstacles = usd_help.get_obstacles_from_stage(
# only_paths=[obstacles_path],
only_paths=["/World"],
reference_prim_path=robot_prim_path,
ignore_substring=[
robot_prim_path,
@@ -319,6 +316,9 @@ def main():
past_orientation = cube_orientation
sim_js = robot.get_joints_state()
if sim_js is None:
print("sim_js is None")
continue
sim_js_names = robot.dof_names
if np.any(np.isnan(sim_js.positions)):
log_error("isaac sim has returned NAN joint position values.")
@@ -362,7 +362,7 @@ def main():
spheres[si].set_radius(float(s.radius))
robot_static = False
if (np.max(np.abs(sim_js.velocities)) < 0.2) or args.reactive:
if (np.max(np.abs(sim_js.velocities)) < 0.5) or args.reactive:
robot_static = True
if (
(

8
examples/isaac_sim/motion_gen_reacher_nvblox.py
View File

@@ -188,8 +188,9 @@ def main():
step_index = my_world.current_time_step_index
# print(step_index)
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -207,6 +208,9 @@ def main():
if target_pose is None:
target_pose = cube_position
sim_js = robot.get_joints_state()
if sim_js is None:
print("sim_js is None")
continue
sim_js_names = robot.dof_names
cu_js = JointState(
position=tensor_args.to_device(sim_js.positions),

21
examples/isaac_sim/mpc_example.py
View File

@@ -116,8 +116,10 @@ def draw_points(rollouts: torch.Tensor):
import random
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
N = 100
# if draw.get_num_points() > 0:
@@ -265,8 +267,10 @@ def main():
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -281,7 +285,7 @@ def main():
if step_index % 1000 == 0:
print("Updating world")
obstacles = usd_help.get_obstacles_from_stage(
# only_paths=[obstacles_path],
only_paths=["/World"],
ignore_substring=[
robot_prim_path,
"/World/target",
@@ -315,6 +319,9 @@ def main():
# get robot current state:
sim_js = robot.get_joints_state()
if sim_js is None:
print("sim_js is None")
continue
js_names = robot.dof_names
sim_js_names = robot.dof_names
@@ -354,7 +361,7 @@ def main():
cmd_state_full = cmd_state
art_action = ArticulationAction(
cmd_state.position.cpu().numpy(),
cmd_state.position.view(-1).cpu().numpy(),
# cmd_state.velocity.cpu().numpy(),
joint_indices=idx_list,
)
@@ -364,7 +371,7 @@ def main():
if succ:
# set desired joint angles obtained from IK:
for _ in range(3):
for _ in range(1):
articulation_controller.apply_action(art_action)
else:

13
examples/isaac_sim/mpc_nvblox_example.py
View File

@@ -104,7 +104,10 @@ def draw_points(rollouts: torch.Tensor):
import random
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
N = 100
@@ -242,8 +245,9 @@ def main():
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -278,6 +282,9 @@ def main():
# get robot current state:
sim_js = robot.get_joints_state()
if sim_js is None:
print("sim_js is None")
continue
js_names = robot.dof_names
sim_js_names = robot.dof_names

14
examples/isaac_sim/multi_arm_reacher.py
View File

@@ -149,6 +149,7 @@ def main():
collision_activation_distance=0.025,
fixed_iters_trajopt=True,
maximum_trajectory_dt=0.5,
ik_opt_iters=500,
)
motion_gen = MotionGen(motion_gen_config)
print("warming up...")
@@ -221,8 +222,9 @@ def main():
step_index = my_world.current_time_step_index
# print(step_index)
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)
@@ -235,7 +237,7 @@ def main():
if step_index == 50 or step_index % 1000 == 0.0:
print("Updating world, reading w.r.t.", robot_prim_path)
obstacles = usd_help.get_obstacles_from_stage(
# only_paths=[obstacles_path],
only_paths=["/World"],
reference_prim_path=robot_prim_path,
ignore_substring=[
robot_prim_path,
@@ -258,6 +260,9 @@ def main():
if target_pose is None:
target_pose = cube_position
sim_js = robot.get_joints_state()
if sim_js is None:
print("sim_js is None")
continue
sim_js_names = robot.dof_names
cu_js = JointState(
position=tensor_args.to_device(sim_js.positions),
@@ -287,11 +292,10 @@ def main():
for si, s in enumerate(sph_list[0]):
spheres[si].set_world_pose(position=np.ravel(s.position))
spheres[si].set_radius(float(s.radius))
# print(sim_js.velocities)
if (
np.linalg.norm(cube_position - target_pose) > 1e-3
and np.linalg.norm(past_pose - cube_position) == 0.0
and np.max(np.abs(sim_js.velocities)) < 0.2
and np.max(np.abs(sim_js.velocities)) < 0.5
):
# Set EE teleop goals, use cube for simple non-vr init:
ee_translation_goal = cube_position

15
examples/isaac_sim/realsense_collision.py
View File

@@ -68,7 +68,10 @@ def draw_points(voxels):
# Third Party
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:
@@ -125,7 +128,10 @@ def clip_camera(camera_data):
def draw_line(start, gradient):
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:
@@ -270,7 +276,10 @@ if __name__ == "__main__":
draw_line(sph_position, d_vec[..., :3].view(3).cpu().numpy())
else:
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:

20
examples/isaac_sim/realsense_mpc.py
View File

@@ -98,7 +98,10 @@ def draw_rollout_points(rollouts: torch.Tensor, clear: bool = False):
import random
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
N = 100
@@ -124,7 +127,10 @@ def draw_points(voxels):
# Third Party
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:
@@ -171,7 +177,10 @@ def clip_camera(camera_data):
def draw_line(start, gradient):
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:
@@ -345,8 +354,9 @@ if __name__ == "__main__":
if cmd_step_idx == 0:
draw_rollout_points(mpc.get_visual_rollouts(), clear=not args.use_debug_draw)
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)

15
examples/isaac_sim/realsense_reacher.py
View File

@@ -88,7 +88,10 @@ def draw_points(voxels):
# Third Party
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:
@@ -135,7 +138,10 @@ def clip_camera(camera_data):
def draw_line(start, gradient):
# Third Party
from omni.isaac.debug_draw import _debug_draw
try:
from omni.isaac.debug_draw import _debug_draw
except ImportError:
from isaacsim.util.debug_draw import _debug_draw
draw = _debug_draw.acquire_debug_draw_interface()
# if draw.get_num_points() > 0:
@@ -283,8 +289,9 @@ if __name__ == "__main__":
continue
step_index = my_world.current_time_step_index
if step_index <= 2:
my_world.reset()
if step_index <= 10:
# my_world.reset()
robot._articulation_view.initialize()
idx_list = [robot.get_dof_index(x) for x in j_names]
robot.set_joint_positions(default_config, idx_list)

4
examples/kinematics_example.py
View File

@@ -8,9 +8,7 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
""" Example computing fk using curobo
"""
"""Example computing fk using curobo"""
# Third Party
import torch

224
examples/mesh_dataset.py Normal file
View File

@@ -0,0 +1,224 @@
#
# 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.
#
import math
import os
import sys
from typing import Optional
import numpy as np
import pyrender
import torch
import torch.nn.functional as F
import trimesh
from torch.utils.data.dataset import Dataset
def fov_and_size_to_intrinsics(fov, img_size, device="cpu"):
img_h, img_w = img_size
fx = img_w / (2 * math.tan(math.radians(fov) / 2))
fy = img_h / (2 * math.tan(math.radians(fov) / 2))
intrinsics = torch.tensor(
[[fx, 0, img_h / 2], [0, fy, img_w / 2], [0, 0, 1]],
dtype=torch.float,
device=device,
)
return intrinsics
def lookat_to_cam_pose(eyes, ats, ups=[[0, 0, 1]], device="cpu", mode="opengl"):
if not isinstance(eyes, torch.Tensor):
eyes = torch.tensor(eyes, device=device, dtype=torch.float32)
if not isinstance(ats, torch.Tensor):
ats = torch.tensor(ats, device=device, dtype=torch.float32)
if not isinstance(ups, torch.Tensor):
ups = torch.tensor(ups, device=device, dtype=torch.float32)
batch_size = eyes.shape[0]
camera_view = F.normalize(ats - eyes, dim=1)
camera_right = F.normalize(torch.cross(camera_view, ups, dim=1), dim=1)
camera_up = F.normalize(torch.cross(camera_right, camera_view, dim=1), dim=1)
# rotation matrix from opencv conventions
T = torch.zeros((batch_size, 4, 4))
if mode == "opengl":
T[:, :3, :] = torch.stack([camera_right, camera_up, -camera_view, eyes], dim=2)
elif mode == "opencv":
T[:, :3, :] = torch.stack([camera_right, -camera_up, camera_view, eyes], dim=2)
else:
raise ValueError(f"Unknown mode: {mode}")
T[:, 3, 3] = 1.0
return T.float()
def sample_sphere_points(N, radius, device="cuda"):
latitude = (torch.rand(size=(N, 1), device=device) - 0.5) * torch.pi
longitude = (torch.rand(size=(N, 1), device=device) - 0.5) * torch.pi * 2
x = torch.cos(latitude) * torch.cos(longitude)
y = torch.sin(latitude) * torch.cos(longitude)
z = torch.sin(longitude)
pc = torch.cat([x, y, z], dim=1) * radius
return pc
def sample_sphere_poses(N, origin, radius, device="cuda"):
eyes = sample_sphere_points(N, radius, device)
if not isinstance(origin, torch.Tensor):
origin = torch.tensor(origin).float().to(device)
ats = origin[None, :].repeat((N, 1))
poses_gl = lookat_to_cam_pose(eyes, ats, device=device, mode="opengl")
poses_cv = lookat_to_cam_pose(eyes, ats, device=device, mode="opencv")
return poses_gl, poses_cv
def compute_origin_and_radius(trimesh_scene):
low, high = trimesh_scene.bounds
center = high + low / 2
low = low - center
high = high - center
radius = max(np.sqrt((high**2).sum()), np.sqrt((low**2).sum()))
return center, radius
def render_batch(trimesh_mesh, camera_poses, fov, image_size):
camera_poses = camera_poses.detach().cpu().numpy()
mesh = pyrender.Mesh.from_trimesh(trimesh_mesh)
scene = pyrender.Scene()
scene.add(mesh)
camera = pyrender.PerspectiveCamera(yfov=fov, aspectRatio=1.0)
camera = pyrender.Node(camera=camera, matrix=np.eye(4))
scene.add_node(camera)
light = pyrender.SpotLight(
color=np.ones(3),
intensity=1.0,
innerConeAngle=np.pi / 16.0,
outerConeAngle=np.pi / 2.0,
)
light = pyrender.Node(light=light, matrix=np.eye(4))
scene.add_node(light)
r = pyrender.OffscreenRenderer(image_size, image_size)
colors = []
depths = []
for camera_pose in camera_poses:
scene.set_pose(camera, camera_pose)
scene.set_pose(light, camera_pose)
color, depth = r.render(scene)
colors.append(color)
depths.append(depth)
return np.asarray(colors), np.asarray(depths)
"""
MeshDataset takes a path to a mesh as input and uses PyRender to render images of the mesh
from a sphere centered around the scene.
"""
class MeshDataset(Dataset):
def __init__(
self,
mesh_file: str = None,
n_frames: int = 10,
image_size: float = 256,
save_data_dir: Optional[str] = None,
trimesh_mesh: Optional[trimesh.Trimesh] = None,
fov_deg: int = 60,
# visible_point: Optional[List[float]] = None,
) -> None:
super().__init__()
self.mesh_file = mesh_file
self.n_frames = n_frames
if trimesh_mesh is None:
self.trimesh_mesh = trimesh.load(self.mesh_file)
else:
self.trimesh_mesh = trimesh_mesh
self.image_size = image_size
origin, radius = compute_origin_and_radius(self.trimesh_mesh)
self.fov_deg = fov_deg
sphere_radius = radius * 2.0
self.camera_poses_gl, self.camera_poses_cv = sample_sphere_poses(
n_frames, origin, sphere_radius, "cuda"
)
self.colors, self.depths = render_batch(
self.trimesh_mesh,
self.camera_poses_gl,
fov=math.radians(self.fov_deg),
image_size=self.image_size,
)
self.intrinsics = fov_and_size_to_intrinsics(
self.fov_deg, (self.image_size, self.image_size), device="cuda"
)
if save_data_dir is not None:
self.save_as_sun3d_dataset(save_data_dir)
# sys.exit(0)
def save_as_sun3d_dataset(self, output_dir):
import imageio
from transforms3d.quaternions import quat2mat
os.makedirs(output_dir, exist_ok=True)
K = self.intrinsics.detach().cpu().numpy().tolist()
intrinsics_text = f"""{K[0][0]} {K[0][1]} {K[0][2]}
{K[1][0]} {K[1][1]} {K[1][2]}
{K[2][0]} {K[2][1]} {K[2][2]}"""
with open(f"{output_dir}/camera-intrinsics.txt", "w") as fp:
fp.write(intrinsics_text)
seqdir = f"{output_dir}/seq-01"
os.makedirs(seqdir, exist_ok=True)
for i in range(len(self)):
data = self[i]
rgb = data["rgba"][:3, :, :].detach().cpu().permute(1, 2, 0).numpy()
depth = data["depth"]
depth = (depth * 1000).detach().cpu().numpy().astype(np.uint16)
nvblox_pose = data["pose"]
eigen_quat = [0.707106769, 0.707106769, 0, 0]
sun3d_to_nvblox_T = torch.eye(4)
sun3d_to_nvblox_T[:3, :3] = torch.tensor(quat2mat(eigen_quat))
sun3d_pose = torch.linalg.inv(sun3d_to_nvblox_T) @ nvblox_pose
P = sun3d_pose.detach().cpu().numpy().tolist()
pose_text = f"""{P[0][0]} {P[0][1]} {P[0][2]} {P[0][3]}
{P[1][0]} {P[1][1]} {P[1][2]} {P[1][3]}
{P[2][0]} {P[2][1]} {P[2][2]} {P[2][3]}
{P[3][0]} {P[3][1]} {P[3][2]} {P[3][3]}"""
framename = f"frame-{str(i).zfill(6)}"
imageio.imwrite(f"{seqdir}/{framename}.color.png", rgb)
imageio.imwrite(f"{seqdir}/{framename}.depth.png", depth)
with open(f"{seqdir}/{framename}.pose.txt", "w") as fp:
fp.write(pose_text)
def __len__(self):
return self.n_frames
def __getitem__(self, index):
rgb_np = self.colors[index]
depth_np = self.depths[index]
a_np = (depth_np > 0).astype(np.uint8) * 255
rgba_np = np.concatenate([rgb_np, a_np[:, :, None]], axis=2)
pose = self.camera_poses_cv[index]
intrinsics = self.intrinsics
depth_np = depth_np.astype(np.float32)
rgba = torch.from_numpy(rgba_np).permute((2, 0, 1))
depth = torch.from_numpy(depth_np).float()
return {"rgba": rgba, "depth": depth, "pose": pose, "intrinsics": intrinsics}

2
examples/pose_sequence_example.py
View File

@@ -8,7 +8,7 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
" This example moves the robot through a sequence of poses and dumps an animated usd."
"This example moves the robot through a sequence of poses and dumps an animated usd."
# CuRobo
from curobo.types.math import Pose
from curobo.types.robot import JointState

67
examples/robot_image_segmentation_example.py
View File

@@ -8,7 +8,7 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
""" This example shows how to use cuRobo's kinematics to generate a mask. """
"""This example shows how to use cuRobo's kinematics to generate a mask."""
# Standard Library
@@ -19,7 +19,7 @@ import imageio
import numpy as np
import torch
import torch.autograd.profiler as profiler
from nvblox_torch.datasets.mesh_dataset import MeshDataset
from mesh_dataset import MeshDataset
from torch.profiler import ProfilerActivity, profile, record_function
# CuRobo
@@ -79,7 +79,7 @@ def create_render_dataset(
mesh_dataset = MeshDataset(
None,
n_frames=n_frames,
image_size=640,
image_size=1920,
save_data_dir=None,
trimesh_mesh=robot_mesh,
fov_deg=fov_deg,
@@ -96,10 +96,14 @@ def mask_image(robot_file="ur5e.yml"):
tensor_args = TensorDeviceType()
curobo_segmenter = RobotSegmenter.from_robot_file(
robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
robot_file,
collision_sphere_buffer=0.01,
distance_threshold=0.05,
use_cuda_graph=True,
ops_dtype=torch.float16,
)
mesh_dataset, q_js = create_render_dataset(robot_file, write_pointcloud, n_frames=20)
mesh_dataset, q_js = create_render_dataset(robot_file, write_pointcloud, n_frames=30)
if save_debug_data:
visualize_scale = 10.0
@@ -179,28 +183,32 @@ def mask_image(robot_file="ur5e.yml"):
)
dt_list = []
for j in range(10):
for i in range(len(mesh_dataset)):
for i in range(len(mesh_dataset)):
data = mesh_dataset[i]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
if not curobo_segmenter.ready:
curobo_segmenter.update_camera_projection(cam_obs)
st_time = time.time()
data = mesh_dataset[i]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs,
q_js,
)
torch.cuda.synchronize()
dt_list.append(time.time() - st_time)
print(
"Segmentation Time (ms), (hz)",
np.mean(dt_list[5:]) * 1000.0,
1.0 / np.mean(dt_list[5:]),
)
if not curobo_segmenter.ready:
curobo_segmenter.update_camera_projection(cam_obs)
st_time = time.time()
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs,
q_js,
)
torch.cuda.synchronize()
dt_list.append(time.time() - st_time)
print("Segmentation Time (ms), (hz)", np.mean(dt_list[5:]) * 1000.0, 1.0 / np.mean(dt_list[5:]))
def batch_mask_image(robot_file="ur5e.yml"):
@@ -216,7 +224,11 @@ def batch_mask_image(robot_file="ur5e.yml"):
tensor_args = TensorDeviceType()
curobo_segmenter = RobotSegmenter.from_robot_file(
robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
robot_file,
collision_sphere_buffer=0.01,
distance_threshold=0.05,
use_cuda_graph=True,
ops_dtype=torch.float16,
)
mesh_dataset, q_js = create_render_dataset(robot_file, save_debug_data, fov_deg=60)
@@ -489,7 +501,8 @@ def profile_mask_image(robot_file="ur5e.yml"):
if __name__ == "__main__":
mask_image()
robot_file = "quad_ur10e.yml" # "ur10e.yml" #"ur5e_robotiq_2f_140.yml"
# mask_image(robot_file)
# profile_mask_image()
# batch_mask_image()
batch_mask_image(robot_file)
# batch_robot_mask_image()