Merge pull request #1413 from awesome-aj0123:usd-integration
PiperOrigin-RevId: 612933036 Change-Id: Ia29230b1d91508dbc8b1cf7c76ff48c17448bf89
This commit is contained in:
Copybara-Service
@@ -0,0 +1,857 @@
|
||||
# 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 typing import List, Optional, Tuple
|
||||
|
||||
import mujoco
|
||||
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 USDMesh:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
model: mujoco.MjModel,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
dataid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
self.stage = stage
|
||||
self.model = model
|
||||
self.geom = geom
|
||||
self.objid = objid
|
||||
self.rgba = rgba
|
||||
self.dataid = dataid
|
||||
self.texture_file = texture_file
|
||||
|
||||
xform_path = f"/World/Mesh_Xform_{objid}"
|
||||
mesh_path = f"{xform_path}/Mesh_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
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()))
|
||||
|
||||
self._attach_material()
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
|
||||
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
|
||||
|
||||
def _attach_material(self):
|
||||
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}")
|
||||
mtl = UsdShade.Material.Define(self.stage, mtl_path)
|
||||
|
||||
if self.texture_file:
|
||||
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"
|
||||
)
|
||||
|
||||
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
||||
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
|
||||
|
||||
# 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)
|
||||
else:
|
||||
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"
|
||||
)
|
||||
|
||||
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
||||
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
|
||||
|
||||
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
|
||||
transformation_mat = mujoco.usd_utils.create_transform_matrix(
|
||||
rotation_matrix=mat, translation_vector=pos
|
||||
).T
|
||||
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
||||
self.update_visibility(visible, frame)
|
||||
|
||||
def update_visibility(self, visible: bool, frame: int):
|
||||
if visible:
|
||||
self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
|
||||
else:
|
||||
self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
|
||||
|
||||
|
||||
class USDPrimitiveMesh:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
self.stage = stage
|
||||
self.geom = geom
|
||||
self.objid = objid
|
||||
self.rgba = rgba
|
||||
self.texture_file = texture_file
|
||||
|
||||
self.usd_prim = Usd.Prim()
|
||||
self.usd_mesh = Usd.Mesh()
|
||||
self.prim_mesh = Usd.PrimMesh()
|
||||
self.transform_op = Usd.TransformOp()
|
||||
|
||||
def _set_refinement_properties(self):
|
||||
self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
|
||||
|
||||
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:
|
||||
x_multiplier, y_multiplier = self.geom.size[:2]
|
||||
|
||||
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)
|
||||
|
||||
def _attach_material(self):
|
||||
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}")
|
||||
mtl = UsdShade.Material.Define(self.stage, mtl_path)
|
||||
if self.texture_file:
|
||||
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"
|
||||
)
|
||||
|
||||
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
||||
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
|
||||
|
||||
# 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)
|
||||
else:
|
||||
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[: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"
|
||||
)
|
||||
|
||||
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
||||
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
|
||||
|
||||
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
|
||||
transformation_mat = mujoco.usd_util.create_transform_matrix(
|
||||
rotation_matrix=mat, translation_vector=pos
|
||||
).T
|
||||
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
||||
self.update_visibility(visible, frame)
|
||||
|
||||
def update_visibility(self, visible: bool, frame: int):
|
||||
if visible:
|
||||
self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
|
||||
else:
|
||||
self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
|
||||
|
||||
|
||||
class USDPrimitive:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
self.stage = stage
|
||||
self.geom = geom
|
||||
self.objid = objid
|
||||
self.rgba = rgba
|
||||
self.texture_file = texture_file
|
||||
|
||||
self.usd_prim = Usd.Prim()
|
||||
self.usd_primitive_shape = Usd.PrimitiveShape()
|
||||
self.transform_op = Usd.TransformOp()
|
||||
|
||||
def _set_refinement_properties(self):
|
||||
self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
|
||||
|
||||
def _attach_material(self):
|
||||
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}")
|
||||
mtl = UsdShade.Material.Define(self.stage, mtl_path)
|
||||
if self.texture_file:
|
||||
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")
|
||||
)
|
||||
|
||||
# settings 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"
|
||||
)
|
||||
|
||||
self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
||||
UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl)
|
||||
|
||||
# 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("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)
|
||||
else:
|
||||
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[: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"
|
||||
)
|
||||
|
||||
self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
|
||||
UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl)
|
||||
|
||||
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
|
||||
transformation_mat = mujoco.usd_util.create_transform_matrix(
|
||||
rotation_matrix=mat, translation_vector=pos
|
||||
).T
|
||||
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
||||
self.update_visibility(visible, frame)
|
||||
|
||||
def update_visibility(self, visible: bool, frame: int):
|
||||
if visible:
|
||||
self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
|
||||
else:
|
||||
self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
|
||||
|
||||
|
||||
class USDCapsule(USDPrimitive):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
|
||||
super().__init__(stage, geom, objid, rgba, texture_file)
|
||||
|
||||
xform_path = f"/World/Capsule_Xform_{objid}"
|
||||
capsule_path = f"{xform_path}/Capsule_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_primitive_shape = UsdGeom.Capsule.Define(stage, capsule_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(capsule_path)
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
self.scale_op = self.usd_xform.AddScaleOp()
|
||||
|
||||
# setting attributes for the shape
|
||||
self._set_size_attributes()
|
||||
self._attach_material()
|
||||
|
||||
self._set_refinement_properties()
|
||||
|
||||
def _set_size_attributes(self):
|
||||
self.usd_primitive_shape.GetRadiusAttr().Set(float(self.geom.size[0]))
|
||||
self.usd_primitive_shape.GetHeightAttr().Set(
|
||||
float(self.geom.size[2] * 2)
|
||||
) # mujoco gives the half length
|
||||
|
||||
|
||||
class USDEllipsoid(USDPrimitive):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
|
||||
super().__init__(stage, geom, objid, rgba, texture_file)
|
||||
|
||||
xform_path = f"/World/Ellipsoid_Xform_{objid}"
|
||||
ellipsoid_path = f"{xform_path}/Ellipsoid_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_primitive_shape = UsdGeom.Sphere.Define(stage, ellipsoid_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(ellipsoid_path)
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
self.scale_op = self.usd_xform.AddScaleOp()
|
||||
|
||||
# setting attributes for the shape
|
||||
self._set_size_attributes()
|
||||
self._attach_material()
|
||||
|
||||
self._set_refinement_properties()
|
||||
|
||||
def _set_size_attributes(self):
|
||||
self.scale_op.Set(Gf.Vec3d(self.geom.size.tolist()))
|
||||
|
||||
|
||||
class USDCubeMesh(USDPrimitiveMesh):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
|
||||
super().__init__(stage, geom, objid, rgba, texture_file)
|
||||
|
||||
xform_path = f"/World/CubeMesh_Xform_{objid}"
|
||||
mesh_path = f"{xform_path}/CubeMesh_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(mesh_path)
|
||||
|
||||
self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
|
||||
width=self.geom.size[0] * 2,
|
||||
height=self.geom.size[1] * 2,
|
||||
depth=self.geom.size[2] * 2,
|
||||
create_uv_map=True,
|
||||
map_texture_to_each_face=True,
|
||||
)
|
||||
|
||||
self.prim_mesh.translate(-self.prim_mesh.get_center())
|
||||
|
||||
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()
|
||||
|
||||
# setting attributes for the shape
|
||||
self._attach_material()
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
|
||||
|
||||
class USDSphereMesh(USDPrimitiveMesh):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
|
||||
super().__init__(stage, geom, objid, rgba, texture_file)
|
||||
|
||||
xform_path = f"/World/SphereMesh_Xform_{objid}"
|
||||
mesh_path = f"{xform_path}/SphereMesh_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(mesh_path)
|
||||
|
||||
self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
|
||||
radius=float(self.geom.size[0]), create_uv_map=True
|
||||
)
|
||||
|
||||
self.prim_mesh.translate(-self.prim_mesh.get_center())
|
||||
|
||||
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()
|
||||
|
||||
# setting attributes for the shape
|
||||
self._attach_material()
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
|
||||
|
||||
class USDCylinderMesh(USDPrimitiveMesh):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
|
||||
super().__init__(stage, geom, objid, rgba, texture_file)
|
||||
|
||||
xform_path = f"/World/CylinderMesh_Xform_{objid}"
|
||||
mesh_path = f"{xform_path}/CylinderMesh_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(mesh_path)
|
||||
|
||||
self.prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
|
||||
radius=self.geom.size[0],
|
||||
height=self.geom.size[2] * 2,
|
||||
create_uv_map=True,
|
||||
)
|
||||
|
||||
self.prim_mesh.translate(-self.prim_mesh.get_center())
|
||||
|
||||
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()
|
||||
|
||||
# setting attributes for the shape
|
||||
self._attach_material()
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
|
||||
|
||||
class USDPlaneMesh(USDPrimitiveMesh):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stage: Usd.Stage,
|
||||
geom: mujoco.MjvGeom,
|
||||
objid: int,
|
||||
rgba: Tuple[int, ...] = (1, 1, 1, 1),
|
||||
texture_file: Optional[str] = None,
|
||||
):
|
||||
|
||||
super().__init__(stage, geom, objid, rgba, texture_file)
|
||||
|
||||
xform_path = f"/World/Plane_Xform_{objid}"
|
||||
plane_path = f"{xform_path}/PlaneMesh_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_mesh = UsdGeom.Mesh.Define(stage, plane_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(plane_path)
|
||||
|
||||
self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
|
||||
width=self.geom.size[0] * 2 if self.geom.size[0] > 0 else 100,
|
||||
height=self.geom.size[1] * 2 if self.geom.size[1] > 0 else 100,
|
||||
depth=0.001,
|
||||
create_uv_map=True,
|
||||
map_texture_to_each_face=True,
|
||||
)
|
||||
|
||||
self.prim_mesh.translate(-self.prim_mesh.get_center())
|
||||
|
||||
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()
|
||||
|
||||
# setting attributes for the shape
|
||||
self._attach_material()
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
|
||||
|
||||
class USDSphereLight:
|
||||
|
||||
def __init__(
|
||||
self, stage: Usd.Stage, objid: int, radius: Optional[float] = 0.3
|
||||
):
|
||||
self.stage = stage
|
||||
|
||||
xform_path = f"/World/Light_Xform_{objid}"
|
||||
light_path = f"{xform_path}/Light_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_light = UsdLux.SphereLight.Define(stage, light_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(light_path)
|
||||
|
||||
# we assume in mujoco that all lights are point lights
|
||||
self.usd_light.GetRadiusAttr().Set(radius)
|
||||
self.usd_light.GetTreatAsPointAttr().Set(False)
|
||||
self.usd_light.GetNormalizeAttr().Set(True)
|
||||
|
||||
# defining ops required by update function
|
||||
self.translate_op = self.usd_xform.AddTranslateOp()
|
||||
|
||||
def update(
|
||||
self, pos: np.ndarray, intensity: int, color: np.ndarray, frame: int
|
||||
):
|
||||
self.translate_op.Set(Gf.Vec3d(pos.tolist()), frame)
|
||||
|
||||
if not np.any(pos):
|
||||
intensity = 0
|
||||
|
||||
self.usd_light.GetIntensityAttr().Set(intensity)
|
||||
self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
|
||||
|
||||
|
||||
class USDDomeLight:
|
||||
|
||||
def __init__(self, stage: Usd.Stage, objid: int):
|
||||
self.stage = stage
|
||||
|
||||
xform_path = f"/World/Light_Xform_{objid}"
|
||||
light_path = f"{xform_path}/Light_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_light = UsdLux.DomeLight.Define(stage, light_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(light_path)
|
||||
|
||||
# we assume in mujoco that all lights are point lights
|
||||
self.usd_light.GetNormalizeAttr().Set(True)
|
||||
|
||||
def update(self, intensity: int, color: np.ndarray, frame: int):
|
||||
|
||||
self.usd_light.GetIntensityAttr().Set(intensity)
|
||||
self.usd_light.GetExposureAttr().Set(0.0)
|
||||
self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
|
||||
|
||||
|
||||
class USDCamera:
|
||||
|
||||
def __init__(self, stage: Usd.Stage, objid: int):
|
||||
self.stage = stage
|
||||
|
||||
xform_path = f"/World/Camera_Xform_{objid}"
|
||||
camera_path = f"{xform_path}/Camera_{objid}"
|
||||
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
|
||||
self.usd_camera = UsdGeom.Camera.Define(stage, camera_path)
|
||||
self.usd_prim = stage.GetPrimAtPath(camera_path)
|
||||
|
||||
# defining ops required by update function
|
||||
self.transform_op = self.usd_xform.AddTransformOp()
|
||||
|
||||
# self.usd_camera.CreateFocalLengthAttr().Set(18.14756) # default in omniverse
|
||||
self.usd_camera.CreateFocalLengthAttr().Set(12)
|
||||
self.usd_camera.CreateFocusDistanceAttr().Set(400)
|
||||
|
||||
self.usd_camera.GetHorizontalApertureAttr().Set(12)
|
||||
|
||||
self.usd_camera.GetClippingRangeAttr().Set(Gf.Vec2f(1e-4, 1e6))
|
||||
|
||||
def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
|
||||
|
||||
transformation_mat = mujoco.usd_util.create_transform_matrix(
|
||||
rotation_matrix=cam_mat, translation_vector=cam_pos
|
||||
).T
|
||||
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
|
||||
@@ -0,0 +1,426 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
import os
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import scipy
|
||||
import termcolor
|
||||
import tqdm
|
||||
|
||||
from typing import List, Optional, Tuple, Union
|
||||
from PIL import Image as im
|
||||
from PIL import ImageOps
|
||||
|
||||
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
|
||||
# pytype: disable=module-attr
|
||||
from pxr import Sdf
|
||||
from pxr import Usd
|
||||
from pxr import UsdGeom
|
||||
|
||||
|
||||
class USDExporter:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: mujoco.MjModel,
|
||||
height: int = 480,
|
||||
width: int = 480,
|
||||
max_geom: int = 10000,
|
||||
output_directory_name: str = "mujoco_usdpkg",
|
||||
output_directory_root: str = "./",
|
||||
light_intensity: int = 10000,
|
||||
camera_names: Optional[List[str]] = None,
|
||||
specialized_materials_file: Optional[str] = None,
|
||||
verbose: bool = True,
|
||||
):
|
||||
"""Initializes a new USD Exporter
|
||||
|
||||
Args:
|
||||
model: an MjModel instance.
|
||||
height: image height in pixels.
|
||||
width: image width in pixels.
|
||||
max_geom: Optional integer specifying the maximum number of geoms that
|
||||
can be rendered in the same scene. If None this will be chosen
|
||||
automatically based on the estimated maximum number of renderable
|
||||
geoms in the model.
|
||||
output_directory_name: name of root directory to store outputted frames
|
||||
and assets generated by the USD renderer.
|
||||
output_directory_root: path to root directory storing generated frames
|
||||
and assets by the USD renderer.
|
||||
verbose: decides whether to print updates.
|
||||
"""
|
||||
|
||||
buffer_width = model.vis.global_.offwidth
|
||||
buffer_height = model.vis.global_.offheight
|
||||
|
||||
if width > buffer_width:
|
||||
raise ValueError(f"""
|
||||
Image width {width} > framebuffer width {buffer_width}. Either reduce the image
|
||||
width or specify a larger offscreen framebuffer in the model XML using the
|
||||
clause:
|
||||
<visual>
|
||||
<global offwidth="my_width"/>
|
||||
</visual>""".lstrip())
|
||||
|
||||
if height > buffer_height:
|
||||
raise ValueError(f"""
|
||||
Image height {height} > framebuffer height {buffer_height}. Either reduce the
|
||||
image height or specify a larger offscreen framebuffer in the model XML using
|
||||
the clause:
|
||||
<visual>
|
||||
<global offheight="my_height"/>
|
||||
</visual>""".lstrip())
|
||||
|
||||
self.model = model
|
||||
self.height = height
|
||||
self.width = width
|
||||
self.max_geom = max_geom
|
||||
self.output_directory_name = output_directory_name
|
||||
self.output_directory_root = output_directory_root
|
||||
self.light_intensity = light_intensity
|
||||
self.camera_names = camera_names
|
||||
self.specialized_materials_file = specialized_materials_file
|
||||
self.verbose = verbose
|
||||
|
||||
self.frame_count = 0 # maintains how many times we have saved the scene
|
||||
self.updates = 0
|
||||
|
||||
self.geom_name2usd = {}
|
||||
|
||||
# initializing rendering requirements
|
||||
self.renderer = mujoco.Renderer(model, height, width, max_geom)
|
||||
self._initialize_usd_stage()
|
||||
self._scene_option = mujoco.MjvOption() # using default scene option
|
||||
|
||||
# initializing output_directories
|
||||
self._initialize_output_directories()
|
||||
|
||||
# loading required textures for the scene
|
||||
self._load_textures()
|
||||
|
||||
@property
|
||||
def usd(self):
|
||||
return self.stage.GetRootLayer().ExportToString()
|
||||
|
||||
@property
|
||||
def scene(self):
|
||||
return self.renderer.scene
|
||||
|
||||
def _initialize_usd_stage(self):
|
||||
self.stage = Usd.Stage.CreateInMemory()
|
||||
UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
|
||||
self.stage.SetStartTimeCode(0)
|
||||
# add as user input
|
||||
self.stage.SetTimeCodesPerSecond(24.0)
|
||||
|
||||
default_prim = UsdGeom.Xform.Define(
|
||||
self.stage, Sdf.Path("/World")
|
||||
).GetPrim()
|
||||
self.stage.SetDefaultPrim(default_prim)
|
||||
|
||||
def _initialize_output_directories(self):
|
||||
self.output_directory_path = os.path.join(
|
||||
self.output_directory_root, self.output_directory_name
|
||||
)
|
||||
if not os.path.exists(self.output_directory_path):
|
||||
os.makedirs(self.output_directory_path)
|
||||
|
||||
self.frames_directory = os.path.join(self.output_directory_path, "frames")
|
||||
if not os.path.exists(self.frames_directory):
|
||||
os.makedirs(self.frames_directory)
|
||||
|
||||
self.assets_directory = os.path.join(self.output_directory_path, "assets")
|
||||
if not os.path.exists(self.assets_directory):
|
||||
os.makedirs(self.assets_directory)
|
||||
|
||||
if self.verbose:
|
||||
print(
|
||||
termcolor.colored(
|
||||
"Writing output frames and assets to"
|
||||
f" {self.output_directory_path}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
|
||||
def update_scene(
|
||||
self,
|
||||
data: mujoco.MjData,
|
||||
scene_option: Optional[mujoco.MjvOption] = None,
|
||||
):
|
||||
"""Updates the scene with latest sim data
|
||||
|
||||
Args:
|
||||
data: structure storing current simulation state
|
||||
scene_option: we use this to determine which geom groups to activate
|
||||
"""
|
||||
|
||||
self.frame_count += 1
|
||||
|
||||
scene_option = scene_option or self._scene_option
|
||||
|
||||
# update the mujoco renderer
|
||||
self.renderer.update_scene(data, scene_option=scene_option)
|
||||
|
||||
# TODO: update scene options
|
||||
if self.updates == 0:
|
||||
self._initialize_usd_stage()
|
||||
|
||||
self._load_lights()
|
||||
self._load_cameras()
|
||||
|
||||
self._update_geoms()
|
||||
self._update_lights()
|
||||
self._update_cameras(data, scene_option=scene_option)
|
||||
|
||||
self.updates += 1
|
||||
|
||||
def _load_textures(self):
|
||||
# TODO: remove code once added internally to mujoco
|
||||
data_adr = 0
|
||||
self.texture_files = []
|
||||
for texture_id in tqdm.tqdm(range(self.model.ntex)):
|
||||
texture_height = self.model.tex_height[texture_id]
|
||||
texture_width = self.model.tex_width[texture_id]
|
||||
pixels = 3 * texture_height * texture_width
|
||||
img = im.fromarray(
|
||||
self.model.tex_rgb[data_adr : data_adr + pixels].reshape(
|
||||
texture_height, texture_width, 3
|
||||
)
|
||||
)
|
||||
img = ImageOps.flip(img)
|
||||
|
||||
texture_file_name = f"texture_{texture_id}.png"
|
||||
|
||||
img.save(os.path.join(self.assets_directory, texture_file_name))
|
||||
|
||||
relative_path = os.path.relpath(
|
||||
self.assets_directory, self.frames_directory
|
||||
)
|
||||
img_path = os.path.join(
|
||||
relative_path, texture_file_name
|
||||
) # relative path, TODO: switch back to this
|
||||
|
||||
self.texture_files.append(img_path)
|
||||
|
||||
data_adr += pixels
|
||||
|
||||
if self.verbose:
|
||||
print(
|
||||
termcolor.colored(
|
||||
f"Completed writing {self.model.ntex} textures to"
|
||||
f" {self.assets_directory}",
|
||||
"green",
|
||||
)
|
||||
)
|
||||
|
||||
def _load_geom(self, geom: mujoco.MjvGeom):
|
||||
|
||||
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
|
||||
assert geom_name not in self.geom_name2usd
|
||||
|
||||
texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None
|
||||
|
||||
# handles meshes in scene
|
||||
if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
|
||||
usd_geom = mujoco.usd.USDMesh(
|
||||
stage=self.stage,
|
||||
model=self.model,
|
||||
geom=geom,
|
||||
objid=geom_name,
|
||||
dataid=self.model.geom_dataid[geom.objid],
|
||||
rgba=geom.rgba,
|
||||
texture_file=texture_file,
|
||||
)
|
||||
elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE:
|
||||
usd_geom = mujoco.usd.USDPlaneMesh(
|
||||
stage=self.stage,
|
||||
geom=geom,
|
||||
objid=geom_name,
|
||||
rgba=geom.rgba,
|
||||
texture_file=texture_file,
|
||||
)
|
||||
elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE:
|
||||
usd_geom = mujoco.usd.USDSphereMesh(
|
||||
stage=self.stage,
|
||||
geom=geom,
|
||||
objid=geom_name,
|
||||
rgba=geom.rgba,
|
||||
texture_file=texture_file,
|
||||
)
|
||||
elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE:
|
||||
usd_geom = mujoco.usd.USDCapsule(
|
||||
stage=self.stage,
|
||||
geom=geom,
|
||||
objid=geom_name,
|
||||
rgba=geom.rgba,
|
||||
texture_file=texture_file,
|
||||
)
|
||||
elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
|
||||
usd_geom = mujoco.usd.USDEllipsoid(
|
||||
stage=self.stage,
|
||||
geom=geom,
|
||||
objid=geom_name,
|
||||
rgba=geom.rgba,
|
||||
texture_file=texture_file,
|
||||
)
|
||||
elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER:
|
||||
usd_geom = mujoco.usd.USDCylinderMesh(
|
||||
stage=self.stage,
|
||||
geom=geom,
|
||||
objid=geom_name,
|
||||
rgba=geom.rgba,
|
||||
texture_file=texture_file,
|
||||
)
|
||||
elif geom.type == mujoco.mjtGeom.mjGEOM_BOX:
|
||||
usd_geom = mujoco.usd.USDCubeMesh(
|
||||
stage=self.stage,
|
||||
geom=geom,
|
||||
objid=geom_name,
|
||||
rgba=geom.rgba,
|
||||
texture_file=texture_file,
|
||||
)
|
||||
else:
|
||||
usd_geom = None
|
||||
|
||||
self.geom_name2usd[geom_name] = usd_geom
|
||||
|
||||
def _update_geoms(self):
|
||||
|
||||
geom_names = set(self.geom_name2usd.keys())
|
||||
|
||||
# iterate through all geoms in the scene and makes update
|
||||
for i in range(self.scene.ngeom):
|
||||
geom = self.scene.geoms[i]
|
||||
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
|
||||
|
||||
if geom_name not in self.geom_name2usd:
|
||||
self._load_geom(geom)
|
||||
if self.geom_name2usd[geom_name]:
|
||||
self.geom_name2usd[geom_name].update_visibility(False, 0)
|
||||
|
||||
if self.geom_name2usd[geom_name]:
|
||||
self.geom_name2usd[geom_name].update(
|
||||
pos=geom.pos,
|
||||
mat=geom.mat,
|
||||
visible=geom.rgba[3] > 0,
|
||||
frame=self.updates,
|
||||
)
|
||||
if geom_name in geom_names:
|
||||
geom_names.remove(geom_name)
|
||||
|
||||
for geom_name in geom_names:
|
||||
if self.geom_name2usd[geom_name]:
|
||||
self.geom_name2usd[geom_name].update_visibility(False, self.updates)
|
||||
|
||||
def _load_lights(self):
|
||||
# initializes an usd light object for every light in the scene
|
||||
self.usd_lights = []
|
||||
for i in range(self.scene.nlight):
|
||||
light = self.scene.lights[i]
|
||||
if not np.allclose(light.pos, [0, 0, 0]):
|
||||
self.usd_lights.append(mujoco.usd.USDSphereLight(stage=self.stage, objid=i))
|
||||
else:
|
||||
self.usd_lights.append(None)
|
||||
|
||||
def _update_lights(self):
|
||||
for i in range(self.scene.nlight):
|
||||
light = self.scene.lights[i]
|
||||
|
||||
if np.allclose(light.pos, [0, 0, 0]):
|
||||
continue
|
||||
|
||||
if self.usd_lights[i] is None:
|
||||
continue
|
||||
|
||||
self.usd_lights[i].update(
|
||||
pos=light.pos,
|
||||
intensity=self.light_intensity,
|
||||
color=light.diffuse,
|
||||
frame=self.updates,
|
||||
)
|
||||
|
||||
print("done updating")
|
||||
|
||||
def _load_cameras(self):
|
||||
self.usd_cameras = []
|
||||
for name in self.camera_names:
|
||||
self.usd_cameras.append(mujoco.usd.USDCamera(stage=self.stage, objid=name))
|
||||
|
||||
def _update_cameras(
|
||||
self,
|
||||
data: mujoco.MjData,
|
||||
scene_option: Optional[mujoco.MjvOption] = None,
|
||||
):
|
||||
for i in range(len(self.usd_cameras)):
|
||||
|
||||
camera = self.usd_cameras[i]
|
||||
camera_name = self.camera_names[i]
|
||||
|
||||
self.renderer.update_scene(
|
||||
data, scene_option=scene_option, camera=camera_name
|
||||
)
|
||||
|
||||
avg_camera = mujoco.mjv_averageCamera(self.scene.camera[0], self.scene.camera[1])
|
||||
|
||||
forward = avg_camera.forward
|
||||
up = avg_camera.up
|
||||
right = np.cross(forward, up)
|
||||
|
||||
R = np.eye(3)
|
||||
R[:, 0] = right
|
||||
R[:, 1] = up
|
||||
R[:, 2] = -forward
|
||||
|
||||
camera.update(cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates)
|
||||
|
||||
def add_light(
|
||||
self,
|
||||
pos: List[float],
|
||||
intensity: int,
|
||||
radius: Optional[float] = 1.0,
|
||||
color: Optional[np.ndarray] = np.array([0.3, 0.3, 0.3]),
|
||||
objid: Optional[int] = 1,
|
||||
light_type: Optional[str] = "sphere",
|
||||
):
|
||||
|
||||
if light_type == "sphere":
|
||||
new_light = mujoco.usd.USDSphereLight(stage=self.stage, objid=objid, radius=radius)
|
||||
|
||||
new_light.update(pos=pos, intensity=intensity, color=color, frame=0)
|
||||
elif light_type == "dome":
|
||||
new_light = mujoco.usd.USDDomeLight(stage=self.stage, objid=objid)
|
||||
|
||||
new_light.update(intensity=intensity, color=color, frame=0)
|
||||
|
||||
def add_camera(
|
||||
self,
|
||||
pos: List[float],
|
||||
rotation_xyz: List[float],
|
||||
objid: Optional[int] = 1,
|
||||
):
|
||||
new_camera = mujoco.usd.USDCamera(stage=self.stage, objid=objid)
|
||||
|
||||
r = scipy.spatial.transform.Rotation.from_euler(
|
||||
"xyz", rotation_xyz, degrees=True)
|
||||
new_camera.update(cam_pos=pos, cam_mat=r.as_matrix(), frame=0)
|
||||
|
||||
def save_scene(self, filetype: str = "usd"):
|
||||
self.stage.SetEndTimeCode(self.frame_count)
|
||||
self.stage.Export(
|
||||
f"{self.output_directory_root}/{self.output_directory_name}/frames/frame_{self.frame_count}_.{filetype}"
|
||||
)
|
||||
if self.verbose:
|
||||
print(termcolor.colored(f"Writing frame_{self.frame_count}", "green"))
|
||||
@@ -0,0 +1,27 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
import numpy as np
|
||||
|
||||
|
||||
def create_transform_matrix(rotation_matrix, translation_vector):
|
||||
# Ensure rotation_matrix and translation_vector are NumPy arrays
|
||||
rotation_matrix = np.array(rotation_matrix)
|
||||
translation_vector = np.array(translation_vector)
|
||||
|
||||
transform_matrix = np.eye(4)
|
||||
transform_matrix[:3, :3] = rotation_matrix
|
||||
transform_matrix[:3, 3] = translation_vector
|
||||
|
||||
return transform_matrix
|
||||
Reference in New Issue
Block a user