# Copyright (c) 2022-2024, The ORBIT Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
# needed to import for allowing type-hinting: np.ndarray | torch.Tensor | None
from __future__ import annotations
import numpy as np
import torch
import trimesh
import warp as wp
from omni.isaac.orbit.utils.warp import raycast_mesh
[docs]def color_meshes_by_height(meshes: list[trimesh.Trimesh], **kwargs) -> trimesh.Trimesh:
"""
Color the vertices of a trimesh object based on the z-coordinate (height) of each vertex,
using the Turbo colormap. If the z-coordinates are all the same, the vertices will be colored
with a single color.
Args:
meshes: A list of trimesh objects.
Keyword Args:
color: A list of 3 integers in the range [0,255] representing the RGB
color of the mesh. Used when the z-coordinates of all vertices are the same.
color_map: The name of the color map to be used. Defaults to "turbo".
Returns:
A trimesh object with the vertices colored based on the z-coordinate (height) of each vertex.
"""
# Combine all meshes into a single mesh
mesh = trimesh.util.concatenate(meshes)
# Get the z-coordinates of each vertex
heights = mesh.vertices[:, 2]
# Check if the z-coordinates are all the same
if np.max(heights) == np.min(heights):
# Obtain a single color: light blue
color = kwargs.pop("color", [172, 216, 230, 255])
color = np.asarray(color, dtype=np.uint8)
# Set the color for all vertices
mesh.visual.vertex_colors = color
else:
# Normalize the heights to [0,1]
heights_normalized = (heights - np.min(heights)) / (np.max(heights) - np.min(heights))
# clip lower and upper bounds to have better color mapping
heights_normalized = np.clip(heights_normalized, 0.1, 0.9)
# Get the color for each vertex based on the height
color_map = kwargs.pop("color_map", "turbo")
colors = trimesh.visual.color.interpolate(heights_normalized, color_map=color_map)
# Set the vertex colors
mesh.visual.vertex_colors = colors
# Return the mesh
return mesh
[docs]def create_prim_from_mesh(prim_path: str, mesh: trimesh.Trimesh, **kwargs):
"""Create a USD prim with mesh defined from vertices and triangles.
The function creates a USD prim with a mesh defined from vertices and triangles. It performs the
following steps:
- Create a USD Xform prim at the path :obj:`prim_path`.
- Create a USD prim with a mesh defined from the input vertices and triangles at the path :obj:`{prim_path}/mesh`.
- Assign a physics material to the mesh at the path :obj:`{prim_path}/physicsMaterial`.
- Assign a visual material to the mesh at the path :obj:`{prim_path}/visualMaterial`.
Args:
prim_path: The path to the primitive to be created.
mesh: The mesh to be used for the primitive.
Keyword Args:
translation: The translation of the terrain. Defaults to None.
orientation: The orientation of the terrain. Defaults to None.
visual_material: The visual material to apply. Defaults to None.
physics_material: The physics material to apply. Defaults to None.
"""
# need to import these here to prevent isaacsim launching when importing this module
import omni.isaac.core.utils.prims as prim_utils
from pxr import UsdGeom
import omni.isaac.orbit.sim as sim_utils
# create parent prim
prim_utils.create_prim(prim_path, "Xform")
# create mesh prim
prim = prim_utils.create_prim(
f"{prim_path}/mesh",
"Mesh",
translation=kwargs.get("translation"),
orientation=kwargs.get("orientation"),
attributes={
"points": mesh.vertices,
"faceVertexIndices": mesh.faces.flatten(),
"faceVertexCounts": np.asarray([3] * len(mesh.faces)),
"subdivisionScheme": "bilinear",
},
)
# apply collider properties
collider_cfg = sim_utils.CollisionPropertiesCfg(collision_enabled=True)
sim_utils.define_collision_properties(prim.GetPrimPath(), collider_cfg)
# add rgba color to the mesh primvars
if mesh.visual.vertex_colors is not None:
# obtain color from the mesh
rgba_colors = np.asarray(mesh.visual.vertex_colors).astype(np.float32) / 255.0
# displayColor is a primvar attribute that is used to color the mesh
color_prim_attr = prim.GetAttribute("primvars:displayColor")
color_prim_var = UsdGeom.Primvar(color_prim_attr)
color_prim_var.SetInterpolation(UsdGeom.Tokens.vertex)
color_prim_attr.Set(rgba_colors[:, :3])
# displayOpacity is a primvar attribute that is used to set the opacity of the mesh
display_prim_attr = prim.GetAttribute("primvars:displayOpacity")
display_prim_var = UsdGeom.Primvar(display_prim_attr)
display_prim_var.SetInterpolation(UsdGeom.Tokens.vertex)
display_prim_var.Set(rgba_colors[:, 3])
# create visual material
if kwargs.get("visual_material") is not None:
visual_material_cfg: sim_utils.VisualMaterialCfg = kwargs.get("visual_material")
# spawn the material
visual_material_cfg.func(f"{prim_path}/visualMaterial", visual_material_cfg)
sim_utils.bind_visual_material(prim.GetPrimPath(), f"{prim_path}/visualMaterial")
# create physics material
if kwargs.get("physics_material") is not None:
physics_material_cfg: sim_utils.RigidBodyMaterialCfg = kwargs.get("physics_material")
# spawn the material
physics_material_cfg.func(f"{prim_path}/physicsMaterial", physics_material_cfg)
sim_utils.bind_physics_material(prim.GetPrimPath(), f"{prim_path}/physicsMaterial")
[docs]def find_flat_patches(
wp_mesh: wp.Mesh,
num_patches: int,
patch_radius: float | list[float],
origin: np.ndarray | torch.Tensor | tuple[float, float, float],
x_range: tuple[float, float],
y_range: tuple[float, float],
z_range: tuple[float, float],
max_height_diff: float,
):
"""Finds flat patches of given radius in the input mesh.
The function finds flat patches of given radius based on the search space defined by the input ranges.
The search space is characterized by origin in the mesh frame, and the x, y, and z ranges. The x and y
ranges are used to sample points in the 2D region around the origin, and the z range is used to filter
patches based on the height of the points.
The function performs rejection sampling to find the patches based on the following steps:
1. Sample patch locations in the 2D region around the origin.
2. Define a ring of points around each patch location to query the height of the points using ray-casting.
3. Reject patches that are outside the z range or have a height difference that is too large.
4. Keep sampling until all patches are valid.
Args:
wp_mesh: The warp mesh to find patches in.
num_patches: The desired number of patches to find.
patch_radius: The radii used to form patches. If a list is provided, multiple patch sizes are checked.
This is useful to deal with holes or other artifacts in the mesh.
origin: The origin defining the center of the search space. This is specified in the mesh frame.
x_range: The range of X coordinates to sample from.
y_range: The range of Y coordinates to sample from.
z_range: The range of valid Z coordinates used for filtering patches.
max_height_diff: The maximum allowable distance between the lowest and highest points
on a patch to consider it as valid. If the difference is greater than this value,
the patch is rejected.
Returns:
A tensor of shape (num_patches, 3) containing the flat patches. The patches are defined in the mesh frame.
Raises:
RuntimeError: If the function fails to find valid patches. This can happen if the input parameters
are not suitable for finding valid patches and maximum number of iterations is reached.
"""
# set device to warp mesh device
device = wp.device_to_torch(wp_mesh.device)
# resolve inputs to consistent type
# -- patch radii
if isinstance(patch_radius, float):
patch_radius = [patch_radius]
# -- origin
if isinstance(origin, np.ndarray):
origin = torch.from_numpy(origin).to(torch.float).to(device)
elif isinstance(origin, torch.Tensor):
origin = origin.to(device)
else:
origin = torch.tensor(origin, dtype=torch.float, device=device)
# create ranges for the x and y coordinates around the origin.
# The provided ranges are bounded by the mesh's bounding box.
x_range = (
max(x_range[0] + origin[0].item(), wp_mesh.points.numpy()[:, 0].min()),
min(x_range[1] + origin[0].item(), wp_mesh.points.numpy()[:, 0].max()),
)
y_range = (
max(y_range[0] + origin[1].item(), wp_mesh.points.numpy()[:, 1].min()),
min(y_range[1] + origin[1].item(), wp_mesh.points.numpy()[:, 1].max()),
)
z_range = (
z_range[0] + origin[2].item(),
z_range[1] + origin[2].item(),
)
# create a circle of points around (0, 0) to query validity of the patches
# the ring of points is uniformly distributed around the circle
angle = torch.linspace(0, 2 * np.pi, 10, device=device)
query_x = []
query_y = []
for radius in patch_radius:
query_x.append(radius * torch.cos(angle))
query_y.append(radius * torch.sin(angle))
query_x = torch.cat(query_x).unsqueeze(1) # dim: (num_radii * 10, 1)
query_y = torch.cat(query_y).unsqueeze(1) # dim: (num_radii * 10, 1)
# dim: (num_radii * 10, 3)
query_points = torch.cat([query_x, query_y, torch.zeros_like(query_x)], dim=-1)
# create buffers
# -- a buffer to store indices of points that are not valid
points_ids = torch.arange(num_patches, device=device)
# -- a buffer to store the flat patches locations
flat_patches = torch.zeros(num_patches, 3, device=device)
# sample points and raycast to find the height.
# 1. Reject points that are outside the z_range or have a height difference that is too large.
# 2. Keep sampling until all points are valid.
iter_count = 0
while len(points_ids) > 0 and iter_count < 10000:
# sample points in the 2D region around the origin
pos_x = torch.empty(len(points_ids), device=device).uniform_(*x_range)
pos_y = torch.empty(len(points_ids), device=device).uniform_(*y_range)
flat_patches[points_ids, :2] = torch.stack([pos_x, pos_y], dim=-1)
# define the query points to check validity of the patch
# dim: (num_patches, num_radii * 10, 3)
points = flat_patches[points_ids].unsqueeze(1) + query_points
points[..., 2] = 100.0
# ray-cast direction is downwards
dirs = torch.zeros_like(points)
dirs[..., 2] = -1.0
# ray-cast to find the height of the patches
ray_hits = raycast_mesh(points.view(-1, 3), dirs.view(-1, 3), wp_mesh)[0]
heights = ray_hits.view(points.shape)[..., 2]
# set the height of the patches
# note: for invalid patches, they would be overwritten in the next iteration
# so it's safe to set the height to the last value
flat_patches[points_ids, 2] = heights[..., -1]
# check validity
# -- height is within the z range
not_valid = torch.any(torch.logical_or(heights < z_range[0], heights > z_range[1]), dim=1)
# -- height difference is within the max height difference
not_valid = torch.logical_or(not_valid, (heights.max(dim=1)[0] - heights.min(dim=1)[0]) > max_height_diff)
# remove invalid patches indices
points_ids = points_ids[not_valid]
# increment count
iter_count += 1
# check all patches are valid
if len(points_ids) > 0:
raise RuntimeError(
"Failed to find valid patches! Please check the input parameters."
f"\n\tMaximum number of iterations reached: {iter_count}"
f"\n\tNumber of invalid patches: {len(points_ids)}"
f"\n\tMaximum height difference: {max_height_diff}"
)
# return the flat patches (in the mesh frame)
return flat_patches - origin
