release repository

tests/nvblox_test.py

@@ -0,0 +1,161 @@
+#
+# 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.geom.sdf.world import CollisionQueryBuffer, WorldCollisionConfig
+from curobo.geom.types import Cuboid, WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.util_file import get_world_configs_path, join_path, load_yaml
+
+try:
+    # Third Party
+    from nvblox_torch.mapper import Mapper
+
+    # CuRobo
+    from curobo.geom.sdf.world_blox import WorldBloxCollision
+
+except ImportError:
+    pytest.skip("Nvblox Torch not available", allow_module_level=True)
+
+
+def test_world_blox_trajectory():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_nvblox.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, -1, 1, 4)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+    d_sph = coll_check.get_swept_sphere_distance(
+        x_sph, query_buffer, weight, act_distance, act_distance + 0.01, 4, True
+    ).view(-1)
+
+    assert d_sph[0] > 10.0
+    assert d_sph[1] == 0.0
+    assert d_sph[2] == 0.0
+
+
+def test_world_blox():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_nvblox.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+    x_sph = torch.as_tensor(
+        [[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
+    ).view(1, 1, -1, 4)
+    # create buffers:
+    query_buffer = CollisionQueryBuffer.initialize_from_shape(
+        x_sph.shape, tensor_args, coll_check.collision_types
+    )
+    weight = tensor_args.to_device([1])
+    act_distance = tensor_args.to_device([0.0])
+    d_sph = coll_check.get_sphere_distance(x_sph, query_buffer, weight, act_distance).view(-1)
+
+    assert d_sph[0] > 10.0
+    assert d_sph[1] == 0.0
+    assert d_sph[2] == 0.0
+
+
+def test_world_blox_bounding():
+    tensor_args = TensorDeviceType()
+    world_file = "collision_nvblox.yml"
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+    bounding_cube = Cuboid("test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1.0, 1.0, 1.0])
+    mesh = coll_check.get_mesh_in_bounding_box(
+        bounding_cube,
+        voxel_size=0.05,
+    )
+    assert len(mesh.vertices) > 0
+
+
+def test_world_blox_get_mesh():
+    world_file = "collision_nvblox.yml"
+    tensor_args = TensorDeviceType()
+
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict)
+    coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
+    coll_check = WorldBloxCollision(coll_cfg)
+
+    world_mesh = coll_check.get_mesh_from_blox_layer("world")
+    assert len(world_mesh.vertices) > 10
+
+
+def test_nvblox_world_from_primitive_world():
+    world_file = "collision_cubby.yml"
+    tensor_args = TensorDeviceType()
+    data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
+    world_cfg = WorldConfig.from_dict(data_dict).get_mesh_world(True)
+    mesh = world_cfg.mesh[0].get_trimesh_mesh()
+    world_cfg.mesh[0].save_as_mesh("world.obj")
+    # Third Party
+    from nvblox_torch.datasets.mesh_dataset import MeshDataset
+
+    # CuRobo
+    from curobo.geom.sdf.world import CollisionCheckerType
+    from curobo.geom.types import BloxMap
+    from curobo.types.camera import CameraObservation
+    from curobo.types.math import Pose
+
+    # create a nvblox collision checker:
+    world_config = WorldConfig(
+        blox=[
+            BloxMap(
+                name="world",
+                pose=[0, 0, 0, 1, 0, 0, 0],
+                voxel_size=0.02,
+                integrator_type="tsdf",
+            )
+        ]
+    )
+    config = WorldCollisionConfig(
+        tensor_args=tensor_args, world_model=world_config, checker_type=CollisionCheckerType.BLOX
+    )
+    world_coll = WorldBloxCollision(config)
+    m_dataset = MeshDataset(
+        None, n_frames=100, image_size=1000, save_data_dir=None, trimesh_mesh=mesh
+    )
+    for i in range(len(m_dataset)):
+        data = m_dataset[i]
+        cam_obs = CameraObservation(
+            rgb_image=data["rgba"],
+            depth_image=data["depth"],
+            intrinsics=data["intrinsics"],
+            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
+        )
+        world_coll.add_camera_frame(cam_obs, "world")
+
+    world_coll.process_camera_frames("world")
+    world_coll.update_blox_mesh("world")
+    integrated_mesh = world_coll.get_mesh_from_blox_layer("world")
+    if len(integrated_mesh.vertices) > 0:
+        assert True
+        # print("saving World")
+        # integrated_mesh.save_as_mesh("collision_test.obj")
+    else:
+        assert True