Abhishek Joshi
498 lines
18 KiB
Python
498 lines
18 KiB
Python
# Copyright 2024 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
from abc import ABC, abstractmethod
|
|
from collections import defaultdict
|
|
from typing import List, Optional, Tuple
|
|
import pprint
|
|
|
|
import mujoco
|
|
|
|
# import mujoco.usd.utils as utils_component
|
|
# import mujoco.usd.shapes as shapes_component
|
|
import utils as utils_component
|
|
import shapes as shapes_component
|
|
|
|
import numpy as np
|
|
|
|
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
|
|
# pytype: disable=module-attr
|
|
from open3d import open3d as o3d
|
|
from pxr import Gf
|
|
from pxr import Sdf
|
|
from pxr import Usd
|
|
from pxr import UsdGeom
|
|
from pxr import UsdLux
|
|
from pxr import UsdShade
|
|
from pxr import Vt
|
|
|
|
class USDObject(ABC):
|
|
""" Abstract interface for all USD objects including meshes and primitives
|
|
|
|
Subclasses must implement:
|
|
|
|
* `_get_uv_geometry(self)`: gets the nessecary UV information to wrap a texture
|
|
around an object in USD. Each subclass implements their own method to getting
|
|
UV information as different objects are contructed in different ways.
|
|
|
|
* `_get_mesh_geometry(self)`: gets the mesh geometry of an object in the scene.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
stage: Usd.Stage,
|
|
geom: mujoco.MjvGeom,
|
|
obj_name: str,
|
|
rgba: np.ndarray = np.array([1, 1, 1, 1]),
|
|
texture_file: Optional[str] = None,
|
|
):
|
|
self.stage = stage
|
|
self.geom = geom
|
|
self.obj_name = obj_name
|
|
self.rgba = rgba
|
|
self.texture_file = texture_file
|
|
|
|
self.xform_path = f"/World/Mesh_Xform_{obj_name}"
|
|
self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
|
|
|
|
# defining ops required by update function
|
|
self.transform_op = self.usd_xform.AddTransformOp()
|
|
self.scale_op = self.usd_xform.AddScaleOp()
|
|
|
|
self.last_visible_frame = -2 # not an arbitary value, forces difference greater than 1 for visibility on 0th frame
|
|
|
|
@abstractmethod
|
|
def _get_uv_geometry(self):
|
|
"""Gets UV information for an object in the scene"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def _get_mesh_geometry(self):
|
|
"""Gets structure of an object in the scene"""
|
|
raise NotImplementedError
|
|
|
|
def attach_image_material(self, usd_mesh):
|
|
"""Attaches an image texture to a material for a USD object"""
|
|
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
|
|
mtl = UsdShade.Material.Define(self.stage, mtl_path)
|
|
|
|
bsdf_shader = UsdShade.Shader.Define(
|
|
self.stage, mtl_path.AppendPath("Principled_BSDF")
|
|
)
|
|
image_shader = UsdShade.Shader.Define(
|
|
self.stage, mtl_path.AppendPath("Image_Texture")
|
|
)
|
|
uvmap_shader = UsdShade.Shader.Define(
|
|
self.stage, mtl_path.AppendPath("uvmap")
|
|
)
|
|
|
|
# setting the bsdf shader attributes
|
|
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
|
|
bsdf_shader.CreateInput(
|
|
"diffuseColor", Sdf.ValueTypeNames.Color3f
|
|
).ConnectToSource(image_shader.ConnectableAPI(), "rgb")
|
|
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
|
|
float(self.rgba[-1])
|
|
)
|
|
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
|
|
self.geom.shininess
|
|
)
|
|
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
|
|
1.0 - self.geom.shininess
|
|
)
|
|
|
|
mtl.CreateSurfaceOutput().ConnectToSource(
|
|
bsdf_shader.ConnectableAPI(), "surface"
|
|
)
|
|
|
|
# setting the image texture attributes
|
|
image_shader.CreateIdAttr("UsdUVTexture")
|
|
image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
|
|
self.texture_file
|
|
)
|
|
image_shader.CreateInput(
|
|
"sourceColorSpace", Sdf.ValueTypeNames.Token
|
|
).Set("sRGB")
|
|
image_shader.CreateInput("wrapS", Sdf.ValueTypeNames.Token).Set("repeat")
|
|
image_shader.CreateInput("wrapT", Sdf.ValueTypeNames.Token).Set("repeat")
|
|
image_shader.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
|
|
uvmap_shader.ConnectableAPI(), "result"
|
|
)
|
|
image_shader.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
|
|
|
|
# setting uvmap shader attributes
|
|
uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2")
|
|
uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap")
|
|
uvmap_shader.CreateOutput("results", Sdf.ValueTypeNames.Float2)
|
|
|
|
mtl.CreateSurfaceOutput().ConnectToSource(
|
|
bsdf_shader.ConnectableAPI(), "surface"
|
|
)
|
|
|
|
usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
|
UsdShade.MaterialBindingAPI(usd_mesh).Bind(mtl)
|
|
|
|
def attach_solid_material(self, usd_mesh):
|
|
"""Attaches an solid texture to a material for a USD object"""
|
|
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
|
|
mtl = UsdShade.Material.Define(self.stage, mtl_path)
|
|
|
|
bsdf_shader = UsdShade.Shader.Define(
|
|
self.stage, mtl_path.AppendPath("Principled_BSDF")
|
|
)
|
|
|
|
# settings the bsdf shader attributes
|
|
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
|
|
|
|
bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(
|
|
tuple(self.rgba[0:3])
|
|
)
|
|
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
|
|
float(self.rgba[-1])
|
|
)
|
|
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
|
|
self.geom.shininess
|
|
)
|
|
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
|
|
1.0 - self.geom.shininess
|
|
)
|
|
|
|
mtl.CreateSurfaceOutput().ConnectToSource(
|
|
bsdf_shader.ConnectableAPI(), "surface"
|
|
)
|
|
|
|
usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
|
UsdShade.MaterialBindingAPI(usd_mesh).Bind(mtl)
|
|
|
|
def _set_refinement_properties(self, usd_prim, scheme="none"):
|
|
usd_prim.GetAttribute("subdivisionScheme").Set(scheme)
|
|
|
|
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int, scale: Optional[np.ndarray] = None):
|
|
"""Updates the position and orientation of an object in the scene for a given frame"""
|
|
transformation_mat = utils_component.create_transform_matrix(
|
|
rotation_matrix=mat, translation_vector=pos
|
|
).T
|
|
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
|
|
|
if visible and frame - self.last_visible_frame > 1:
|
|
# non consecutive visible frames
|
|
self.update_visibility(False, max(0, self.last_visible_frame))
|
|
self.update_visibility(True, frame)
|
|
|
|
if visible:
|
|
self.last_visible_frame = frame
|
|
|
|
if scale is not None:
|
|
self.update_scale(scale, frame)
|
|
|
|
def update_visibility(self, visible: bool, frame: int):
|
|
"""Updates the visibility of an object in a scene for a given frame"""
|
|
visibility_setting = "inherited" if visible else "invisible"
|
|
self.usd_xform.GetVisibilityAttr().Set(visibility_setting, frame)
|
|
|
|
def update_scale(self, scale: np.ndarray, frame: int):
|
|
"""Updates the scale of an object in the scene for a given frame"""
|
|
self.scale_op.Set(Gf.Vec3f(scale.tolist()), frame)
|
|
|
|
class USDMesh(USDObject):
|
|
|
|
def __init__(
|
|
self,
|
|
stage: Usd.Stage,
|
|
model: mujoco.MjModel,
|
|
geom: mujoco.MjvGeom,
|
|
obj_name: str,
|
|
dataid: int,
|
|
rgba: np.ndarray = np.array([1, 1, 1, 1]),
|
|
texture_file: Optional[str] = None,
|
|
):
|
|
super().__init__(stage,
|
|
geom,
|
|
obj_name,
|
|
rgba,
|
|
texture_file)
|
|
|
|
self.model = model
|
|
self.dataid = dataid
|
|
|
|
mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
|
|
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
|
|
self.usd_prim = stage.GetPrimAtPath(mesh_path)
|
|
|
|
# setting mesh structure properties
|
|
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
|
|
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
|
|
self.usd_mesh.GetFaceVertexCountsAttr().Set(
|
|
[3 for _ in range(mesh_facenum)]
|
|
)
|
|
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
|
|
|
|
# setting mesh uv properties
|
|
mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
|
|
self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
|
|
"UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
|
|
)
|
|
self.texcoords.Set(mesh_texcoord)
|
|
self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist()))
|
|
|
|
if self.texture_file:
|
|
self.attach_image_material(self.usd_mesh)
|
|
else:
|
|
self.attach_solid_material(self.usd_mesh)
|
|
|
|
def get_facetexcoord_ranges(self, nmesh, arr):
|
|
facetexcoords_ranges = [0]
|
|
running_sum = 0
|
|
for i in range(nmesh):
|
|
running_sum += arr[i] * 3
|
|
facetexcoords_ranges.append(running_sum)
|
|
return facetexcoords_ranges
|
|
|
|
def _get_uv_geometry(self):
|
|
mesh_texcoord_adr_from = self.model.mesh_texcoordadr[self.dataid]
|
|
mesh_texcoord_adr_to = (
|
|
self.model.mesh_texcoordadr[self.dataid + 1]
|
|
if self.dataid < self.model.nmesh - 1
|
|
else len(self.model.mesh_texcoord)
|
|
)
|
|
mesh_texcoord = self.model.mesh_texcoord[
|
|
mesh_texcoord_adr_from:mesh_texcoord_adr_to
|
|
]
|
|
|
|
mesh_facetexcoord_ranges = self.get_facetexcoord_ranges(
|
|
self.model.nmesh, self.model.mesh_facenum
|
|
)
|
|
|
|
mesh_facetexcoord = self.model.mesh_facetexcoord.flatten()
|
|
mesh_facetexcoord = mesh_facetexcoord[
|
|
mesh_facetexcoord_ranges[self.dataid] : mesh_facetexcoord_ranges[
|
|
self.dataid + 1
|
|
]
|
|
]
|
|
|
|
mesh_facetexcoord[mesh_facetexcoord == len(mesh_texcoord)] = 0
|
|
return mesh_texcoord, mesh_facetexcoord
|
|
|
|
def _get_mesh_geometry(self):
|
|
mesh_vert_adr_from = self.model.mesh_vertadr[self.dataid]
|
|
mesh_vert_adr_to = (
|
|
self.model.mesh_vertadr[self.dataid + 1]
|
|
if self.dataid < self.model.nmesh - 1
|
|
else len(self.model.mesh_vert)
|
|
)
|
|
mesh_vert = self.model.mesh_vert[mesh_vert_adr_from:mesh_vert_adr_to]
|
|
|
|
mesh_face_adr_from = self.model.mesh_faceadr[self.dataid]
|
|
mesh_face_adr_to = (
|
|
self.model.mesh_faceadr[self.dataid + 1]
|
|
if self.dataid < self.model.nmesh - 1
|
|
else len(self.model.mesh_face)
|
|
)
|
|
mesh_face = self.model.mesh_face[mesh_face_adr_from:mesh_face_adr_to]
|
|
mesh_facenum = self.model.mesh_facenum[self.dataid]
|
|
return mesh_vert, mesh_face, mesh_facenum
|
|
|
|
class USDPrimitiveMesh(USDObject):
|
|
|
|
def __init__(
|
|
self,
|
|
mesh_config: dict,
|
|
stage: Usd.Stage,
|
|
geom: mujoco.MjvGeom,
|
|
obj_name: str,
|
|
rgba: np.ndarray = np.array([1, 1, 1, 1]),
|
|
texture_file: Optional[str] = None,
|
|
):
|
|
super().__init__(stage,
|
|
geom,
|
|
obj_name,
|
|
rgba,
|
|
texture_file)
|
|
|
|
self.mesh_config = mesh_config
|
|
self.prim_mesh = self.generate_primitive_mesh()
|
|
|
|
mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
|
|
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
|
|
self.usd_prim = stage.GetPrimAtPath(mesh_path)
|
|
|
|
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
|
|
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
|
|
self.usd_mesh.GetFaceVertexCountsAttr().Set(
|
|
[3 for _ in range(mesh_facenum)]
|
|
)
|
|
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
|
|
|
|
# setting mesh uv properties
|
|
mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
|
|
self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
|
|
"UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
|
|
)
|
|
|
|
self.texcoords.Set(mesh_texcoord)
|
|
self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum * 3)]))
|
|
|
|
self._set_refinement_properties(self.usd_prim)
|
|
|
|
if self.texture_file:
|
|
self.attach_image_material(self.usd_mesh)
|
|
else:
|
|
self.attach_solid_material(self.usd_mesh)
|
|
|
|
def generate_primitive_mesh(self):
|
|
_, prim_mesh = shapes_component.mesh_generator(self.mesh_config)
|
|
prim_mesh.translate(-prim_mesh.get_center())
|
|
return prim_mesh
|
|
|
|
def _get_uv_geometry(self):
|
|
|
|
assert self.prim_mesh
|
|
|
|
x_scale, y_scale = self.geom.texrepeat
|
|
mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
|
|
mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
|
|
|
|
x_multiplier, y_multiplier = 1, 1
|
|
if self.geom.texuniform:
|
|
if "box" in self.mesh_config:
|
|
x_multiplier, y_multiplier = self.mesh_config["box"]["width"], self.mesh_config["box"]["height"]
|
|
elif "sphere" in self.mesh_config:
|
|
x_multiplier, y_multiplier = self.mesh_config["sphere"]["radius"], self.mesh_config["sphere"]["radius"]
|
|
|
|
mesh_texcoord[:, 0] *= x_scale * x_multiplier
|
|
mesh_texcoord[:, 1] *= y_scale * y_multiplier
|
|
|
|
return mesh_texcoord, mesh_facetexcoord.flatten()
|
|
|
|
def _get_mesh_geometry(self):
|
|
assert self.prim_mesh
|
|
|
|
# get mesh geometry from the open3d mesh model
|
|
mesh_vert = np.asarray(self.prim_mesh.vertices)
|
|
mesh_face = np.asarray(self.prim_mesh.triangles)
|
|
|
|
return mesh_vert, mesh_face, len(mesh_face)
|
|
|
|
class USDTendon(USDObject):
|
|
|
|
def __init__(
|
|
self,
|
|
mesh_config: dict,
|
|
stage: Usd.Stage,
|
|
geom: mujoco.MjvGeom,
|
|
obj_name: str,
|
|
rgba: np.ndarray = np.array([1, 1, 1, 1]),
|
|
texture_file: Optional[str] = None,
|
|
):
|
|
super().__init__(stage,
|
|
geom,
|
|
obj_name,
|
|
rgba,
|
|
texture_file)
|
|
|
|
self.mesh_config = mesh_config
|
|
self.tendon_parts = self.generate_primitive_mesh()
|
|
self.usd_refs = defaultdict(dict)
|
|
|
|
for name, mesh in self.tendon_parts.items():
|
|
part_xform_path = f"{self.xform_path}/Mesh_Xform_{name}"
|
|
mesh_path = f"{part_xform_path}/Mesh_{obj_name}"
|
|
self.usd_refs[name]["usd_xform"] = UsdGeom.Xform.Define(stage, part_xform_path)
|
|
self.usd_refs[name]["usd_mesh"] = UsdGeom.Mesh.Define(stage, mesh_path)
|
|
self.usd_refs[name]["usd_prim"] = stage.GetPrimAtPath(mesh_path)
|
|
# adding ops for each of the part xforms
|
|
self.usd_refs[name]["translate_op"] = self.usd_refs[name]["usd_xform"].AddTranslateOp()
|
|
self.usd_refs[name]["scale_op"] = self.usd_refs[name]["usd_xform"].AddScaleOp()
|
|
|
|
# setting mesh geometry properties for each of the parts in the tendon
|
|
part_geometries = self._get_mesh_geometry()
|
|
for name, part_geometry in part_geometries.items():
|
|
self.usd_refs[name]["usd_mesh"].GetPointsAttr().Set(part_geometry["mesh_vert"])
|
|
self.usd_refs[name]["usd_mesh"].GetFaceVertexCountsAttr().Set(
|
|
[3 for _ in range(part_geometry["mesh_facenum"])]
|
|
)
|
|
self.usd_refs[name]["usd_mesh"].GetFaceVertexIndicesAttr().Set(part_geometry["mesh_face"])
|
|
|
|
# setting uv properties for each of the parts in the tendon
|
|
part_uv_geometries = self._get_uv_geometry()
|
|
for name, part_uv_geometry in part_uv_geometries.items():
|
|
self.texcoords = UsdGeom.PrimvarsAPI(self.usd_refs[name]["usd_mesh"]).CreatePrimvar(
|
|
"UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
|
|
)
|
|
self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
|
|
self.texcoords.SetIndices(Vt.IntArray([i for i in range(part_geometry["mesh_facenum"] * 3)]))
|
|
|
|
for name in self.usd_refs.keys():
|
|
self._set_refinement_properties(self.usd_refs[name]["usd_prim"])
|
|
if self.texture_file:
|
|
self.attach_image_material(self.usd_refs[name]["usd_mesh"])
|
|
else:
|
|
self.attach_solid_material(self.usd_refs[name]["usd_mesh"])
|
|
|
|
def generate_primitive_mesh(self):
|
|
mesh_parts = {}
|
|
for part_config in self.mesh_config:
|
|
mesh_name, prim_mesh = shapes_component.mesh_generator(part_config)
|
|
prim_mesh.translate(-prim_mesh.get_center())
|
|
mesh_parts[mesh_name] = prim_mesh
|
|
return mesh_parts
|
|
|
|
def _get_uv_geometry(self):
|
|
part_uv_geometries = defaultdict(dict)
|
|
for name, mesh in self.tendon_parts.items():
|
|
x_scale, y_scale = self.geom.texrepeat
|
|
mesh_texcoord = np.array(mesh.triangle_uvs)
|
|
mesh_facetexcoord = np.asarray(mesh.triangles)
|
|
part_uv_geometries[name] = {
|
|
"mesh_texcoord": mesh_texcoord,
|
|
"mesh_facetexcoord": mesh_facetexcoord
|
|
}
|
|
return part_uv_geometries
|
|
|
|
def _get_mesh_geometry(self):
|
|
part_geometries = defaultdict(dict)
|
|
for name, mesh in self.tendon_parts.items():
|
|
# get mesh geometry from the open3d mesh model
|
|
mesh_vert = np.asarray(mesh.vertices)
|
|
mesh_face = np.asarray(mesh.triangles)
|
|
part_geometries[name] = {
|
|
"mesh_vert": mesh_vert,
|
|
"mesh_face": mesh_face,
|
|
"mesh_facenum": len(mesh_face)
|
|
}
|
|
return part_geometries
|
|
|
|
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int, scale: Optional[np.ndarray] = None):
|
|
"""Updates the position and orientation of an object in the scene for a given frame"""
|
|
super().update(pos, mat, visible, frame, scale)
|
|
for name in self.tendon_parts.keys():
|
|
if "left" in name:
|
|
translate = [0, 0, -scale[2] - (scale[0] / 2)]
|
|
self.usd_refs[name]["translate_op"].Set(Gf.Vec3f(translate), frame)
|
|
elif "right" in name:
|
|
translate = [0, 0, scale[2] + (scale[0] / 2)]
|
|
self.usd_refs[name]["translate_op"].Set(Gf.Vec3f(translate), frame)
|
|
|
|
def update_scale(self, scale: np.ndarray, frame: int):
|
|
"""Updates the scale of the tendon"""
|
|
for name in self.tendon_parts.keys():
|
|
if "cylinder" in name:
|
|
self.usd_refs[name]["scale_op"].Set(Gf.Vec3f(scale.tolist()), frame)
|
|
else:
|
|
hemisphere_scale = scale.tolist()
|
|
hemisphere_scale[2] = hemisphere_scale[0]
|
|
self.usd_refs[name]["scale_op"].Set(Gf.Vec3f(hemisphere_scale), frame)
|
|
|
|
|