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

tests/robot_segmentation_test.py

@@ -0,0 +1,209 @@
+#
+# 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 pytest
+import torch
+
+# CuRobo
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
+from curobo.geom.types import PointCloud, WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.camera import CameraObservation
+from curobo.types.math import Pose
+from curobo.types.robot import RobotConfig
+from curobo.types.state import JointState
+from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
+from curobo.wrap.model.robot_segmenter import RobotSegmenter
+
+torch.manual_seed(30)
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+try:
+    # Third Party
+    from nvblox_torch.datasets.mesh_dataset import MeshDataset
+
+except ImportError:
+    pytest.skip(
+        "Nvblox Torch is not available or pyrender is not installed", allow_module_level=True
+    )
+
+
+def create_render_dataset(
+    robot_file,
+    fov_deg: float = 60,
+    n_frames: int = 20,
+    retract_delta: float = 0.0,
+):
+    # load robot:
+    robot_dict = load_yaml(join_path(get_robot_configs_path(), robot_file))
+    robot_dict["robot_cfg"]["kinematics"]["load_link_names_with_mesh"] = True
+    robot_dict["robot_cfg"]["kinematics"]["load_meshes"] = True
+
+    robot_cfg = RobotConfig.from_dict(robot_dict["robot_cfg"])
+
+    kin_model = CudaRobotModel(robot_cfg.kinematics)
+
+    q = kin_model.retract_config
+
+    q += retract_delta
+
+    meshes = kin_model.get_robot_as_mesh(q)
+
+    world = WorldConfig(mesh=meshes[:])
+    world_table = WorldConfig.from_dict(
+        load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
+    )
+    world_table.cuboid[0].dims = [0.5, 0.5, 0.1]
+    world.add_obstacle(world_table.objects[0])
+    robot_mesh = (
+        WorldConfig.create_merged_mesh_world(world, process_color=False).mesh[0].get_trimesh_mesh()
+    )
+
+    mesh_dataset = MeshDataset(
+        None,
+        n_frames=n_frames,
+        image_size=480,
+        save_data_dir=None,
+        trimesh_mesh=robot_mesh,
+        fov_deg=fov_deg,
+    )
+    q_js = JointState(position=q, joint_names=kin_model.joint_names)
+
+    return mesh_dataset, q_js
+
+
+@pytest.mark.parametrize(
+    "robot_file",
+    ["iiwa.yml", "iiwa_allegro.yml", "franka.yml", "ur10e.yml", "ur5e.yml"],
+)
+def test_mask_image(robot_file):
+    # create robot segmenter:
+    tensor_args = TensorDeviceType()
+
+    curobo_segmenter = RobotSegmenter.from_robot_file(
+        robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
+    )
+
+    mesh_dataset, q_js = create_render_dataset(robot_file, n_frames=5)
+
+    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)
+
+        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
+            cam_obs,
+            q_js,
+        )
+        if torch.count_nonzero(depth_mask) > 100:
+            return
+    assert False
+
+
+@pytest.mark.parametrize(
+    "robot_file",
+    ["iiwa.yml", "iiwa_allegro.yml", "franka.yml", "ur10e.yml", "ur5e.yml"],
+)
+def test_batch_mask_image(robot_file):
+    # create robot segmenter:
+    tensor_args = TensorDeviceType()
+
+    curobo_segmenter = RobotSegmenter.from_robot_file(
+        robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
+    )
+
+    mesh_dataset, q_js = create_render_dataset(robot_file, n_frames=5)
+    mesh_dataset_zoom, q_js = create_render_dataset(robot_file, fov_deg=40, n_frames=5)
+
+    for i in range(len(mesh_dataset)):
+
+        data = mesh_dataset[i]
+        data_zoom = mesh_dataset_zoom[i]
+
+        cam_obs = CameraObservation(
+            depth_image=tensor_args.to_device(data["depth"]) * 1000,
+            intrinsics=data["intrinsics"],
+            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+        )
+        cam_obs_zoom = CameraObservation(
+            depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
+            intrinsics=data_zoom["intrinsics"],
+            pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
+        )
+        cam_obs = cam_obs.stack(cam_obs_zoom)
+        if not curobo_segmenter.ready:
+            curobo_segmenter.update_camera_projection(cam_obs)
+
+        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
+            cam_obs,
+            q_js,
+        )
+        if torch.count_nonzero(depth_mask[0]) > 100 and (torch.count_nonzero(depth_mask[1]) > 100):
+            return
+    assert False
+
+
+@pytest.mark.parametrize(
+    "robot_file",
+    ["iiwa.yml", "iiwa_allegro.yml", "franka.yml", "ur10e.yml", "ur5e.yml"],
+)
+def test_batch_robot_mask_image(robot_file):
+    # create robot segmenter:
+    tensor_args = TensorDeviceType()
+
+    curobo_segmenter = RobotSegmenter.from_robot_file(
+        robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
+    )
+
+    mesh_dataset, q_js = create_render_dataset(robot_file, n_frames=5)
+    mesh_dataset_zoom, q_js_zoom = create_render_dataset(
+        robot_file,
+        fov_deg=40,
+        n_frames=5,
+        retract_delta=0.4,
+    )
+    q_js = q_js.unsqueeze(0).stack(q_js_zoom.unsqueeze(0))
+
+    for i in range(len(mesh_dataset)):
+
+        data = mesh_dataset[i]
+        data_zoom = mesh_dataset_zoom[i]
+
+        cam_obs = CameraObservation(
+            depth_image=tensor_args.to_device(data["depth"]) * 1000,
+            intrinsics=data["intrinsics"],
+            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+        )
+        cam_obs_zoom = CameraObservation(
+            depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
+            intrinsics=data_zoom["intrinsics"],
+            pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
+        )
+        cam_obs = cam_obs.stack(cam_obs_zoom)
+        if not curobo_segmenter.ready:
+            curobo_segmenter.update_camera_projection(cam_obs)
+
+        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
+            cam_obs,
+            q_js,
+        )
+        if torch.count_nonzero(depth_mask[0]) > 100 and (torch.count_nonzero(depth_mask[1]) > 100):
+            return
+    assert False