135 lines
4.6 KiB
Python
135 lines
4.6 KiB
Python
#
|
|
# 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.
|
|
#
|
|
|
|
# Standard Library
|
|
from typing import Dict, List
|
|
|
|
# Third Party
|
|
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
|
|
|
|
try:
|
|
# Third Party
|
|
from omni.isaac.urdf import _urdf # isaacsim 2022.2
|
|
except ImportError:
|
|
from omni.importer.urdf import _urdf # isaac sim 2023.1
|
|
# CuRobo
|
|
from curobo.util_file import get_assets_path, get_filename, get_path_of_dir, join_path
|
|
|
|
|
|
############################################################
|
|
def add_robot_to_scene(
|
|
robot_config: Dict,
|
|
my_world: World,
|
|
load_from_usd: bool = False,
|
|
subroot: str = "",
|
|
robot_name: str = "robot",
|
|
position: np.array = np.array([0, 0, 0]),
|
|
):
|
|
urdf_interface = _urdf.acquire_urdf_interface()
|
|
|
|
import_config = _urdf.ImportConfig()
|
|
import_config.merge_fixed_joints = False
|
|
import_config.convex_decomp = False
|
|
import_config.import_inertia_tensor = True
|
|
import_config.fix_base = True
|
|
import_config.make_default_prim = False
|
|
import_config.self_collision = False
|
|
import_config.create_physics_scene = True
|
|
import_config.import_inertia_tensor = False
|
|
import_config.default_drive_strength = 20000
|
|
import_config.default_position_drive_damping = 500
|
|
import_config.default_drive_type = _urdf.UrdfJointTargetType.JOINT_DRIVE_POSITION
|
|
import_config.distance_scale = 1
|
|
import_config.density = 0.0
|
|
full_path = join_path(get_assets_path(), robot_config["kinematics"]["urdf_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,
|
|
)
|
|
|
|
# prim_path = omni.usd.get_stage_next_free_path(
|
|
# my_world.scene.stage, str(my_world.scene.stage.GetDefaultPrim().GetPath()) + robot_path, False)
|
|
# print(prim_path)
|
|
# robot_prim = my_world.scene.stage.OverridePrim(prim_path)
|
|
# robot_prim.GetReferences().AddReference(dest_path)
|
|
|
|
robot = my_world.scene.add(
|
|
Robot(
|
|
prim_path=robot_path,
|
|
name=robot_name,
|
|
position=position,
|
|
)
|
|
)
|
|
|
|
return robot, robot_path
|
|
|
|
|
|
class VoxelManager:
|
|
def __init__(
|
|
self,
|
|
num_voxels: int = 5000,
|
|
size: float = 0.02,
|
|
color: List[float] = [1, 1, 1],
|
|
prefix_path: str = "/World/curobo/voxel_",
|
|
material_path: str = "/World/looks/v_",
|
|
) -> None:
|
|
self.cuboid_list = []
|
|
self.cuboid_material_list = []
|
|
self.disable_idx = num_voxels
|
|
for i in range(num_voxels):
|
|
target_material = OmniPBR("/World/looks/v_" + str(i), color=np.ravel(color))
|
|
|
|
cube = cuboid.VisualCuboid(
|
|
prefix_path + str(i),
|
|
position=np.array([0, 0, -10]),
|
|
orientation=np.array([1, 0, 0, 0]),
|
|
size=size,
|
|
visual_material=target_material,
|
|
)
|
|
self.cuboid_list.append(cube)
|
|
self.cuboid_material_list.append(target_material)
|
|
cube.set_visibility(True)
|
|
|
|
def update_voxels(self, voxel_position: np.ndarray, color_axis: int = 0):
|
|
max_index = min(voxel_position.shape[0], len(self.cuboid_list))
|
|
|
|
jet = cm.get_cmap("hot") # .reversed()
|
|
z_val = voxel_position[:, 0]
|
|
|
|
jet_colors = jet(z_val)
|
|
|
|
for i in range(max_index):
|
|
self.cuboid_list[i].set_visibility(True)
|
|
|
|
self.cuboid_list[i].set_local_pose(translation=voxel_position[i])
|
|
self.cuboid_material_list[i].set_color(jet_colors[i][:3])
|
|
|
|
for i in range(max_index, len(self.cuboid_list)):
|
|
self.cuboid_list[i].set_local_pose(translation=np.ravel([0, 0, -10.0]))
|
|
|
|
# self.cuboid_list[i].set_visibility(False)
|
|
|
|
def clear(self):
|
|
for i in range(len(self.cuboid_list)):
|
|
self.cuboid_list[i].set_local_pose(translation=np.ravel([0, 0, -10.0]))
|