chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adding mujoco uds bridge

Browse Source
This commit is contained in:
Abhishek Joshi
2024-02-01 02:47:45 -06:00
parent 08bdc9f691
commit 75c0647bcc
5 changed files with 926 additions and 597 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dist/mujoco-3.0.2.tar.gz
Vendored
BIN
View File
Binary file not shown.
python/mujoco/usd_component.py
+581 -356
View File
@@ -1,407 +1,632 @@
import os import numpy as np
from enum import Enum import open3d as o3d
from usd_utils import *
from mujoco import mjtGeom
from pxr import Gf, Sdf, Vt
from typing import Optional, List
from mujoco import _structs, _enums
from pxr import Usd, UsdGeom, UsdShade, UsdLux
import mujoco class USDMesh:
from mujoco.usd_utilities import *
from pxr import Usd, UsdGeom, UsdLux, UsdShade, Vt, Gf, Sdf
from scipy.spatial.transform import Rotation as R
# TODO: clean this up and remove the if statements def __init__(
def create_usd_geom_primitive(geom, stage, texture_file): self,
geom_type = geom.type stage: Usd.Stage,
if geom_type==USDGeomType.Plane.value: model: _structs.MjModel,
return USDPlane(geom, stage, texture_file) geom: _structs.MjvGeom,
elif geom_type==USDGeomType.Sphere.value: objid: int,
return USDSphere(geom, stage, texture_file) dataid: int,
elif geom_type==USDGeomType.Capsule.value: rgba: List[int] = [1,1,1,1],
return USDCapsule(geom, stage, texture_file) texture_file: Optional[str] = None
elif geom_type==USDGeomType.Cylinder.value: ):
return USDCylinder(geom, stage, texture_file) """ Initializes a new USD mesh
elif geom_type==USDGeomType.Cube.value: Args:
return USDCube(geom, stage, texture_file) model: an MjModel instance.
elif geom_type==USDGeomType.Cube.value: dataid: id of the mesh
return USDCube(geom, stage, texture_file) texture_file: texture associated with the mesh
else:
return None
class USDGeomType(Enum):
""" """
Represents different types of geoms we can add to USD
The values match those found by the enum presented here:
https://mujoco.readthedocs.io/en/latest/APIreference/APItypes.html#mjtgeom
"""
Plane = 0
# Hfield = 1
Sphere = 2
Capsule = 3
# Ellipsoid = 4
Cylinder = 5
Cube = 6
Mesh = 7
class USDGeom(object):
"""
Parent class for created geoms
"""
def __init__(self,
geom=None,
stage=None,
texture_file=None):
self.geom = geom
self.stage = stage self.stage = stage
self.model = model
self.geom = geom
self.objid = objid
self.rgba = rgba
self.dataid = dataid
self.texture_file = texture_file self.texture_file = texture_file
self.type = None
self.xform = None
self.prim = None
self.ref = None
# xform operations, set when updating geoms xform_path = f'/World/Mesh_Xform_{objid}'
self.translate_op = None mesh_path= f'{xform_path}/Mesh_{objid}'
self.rotate_op = None self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.scale_op = None self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path)
def set_texture(self): # setting mesh structure properties
if self.texture_file: mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(self.texture_file))[0]}") 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]]
return mesh_texcoord, mesh_facetexcoord
def _get_mesh_geometry(self):
# get mesh geometry structure from reading the mjModel
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) mtl = UsdShade.Material.Define(self.stage, mtl_path)
shader = UsdShade.Shader.Define(self.stage, mtl_path.AppendPath("Shader"))
shader.CreateIdAttr("UsdPreviewSurface")
shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0))
shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
diffuse_tx = UsdShade.Shader.Define(self.stage, mtl_path.AppendPath("DiffuseColorTx")) if self.texture_file:
diffuse_tx.CreateIdAttr('UsdUVTexture') 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"))
diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(self.texture_file) # setting the bsdf shader attributes
diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3) bsdf_shader.CreateIdAttr("UsdPreviewSurface")
shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb') bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(image_shader.ConnectableAPI(), "rgb")
mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface") bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(float(self.rgba[-1]))
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
UsdShade.MaterialBindingAPI(self.prim).Bind(mtl)
def update_geom(self, new_geom): self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
self.update_pos(new_geom.pos) UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
self.update_rotation(new_geom.mat)
self.update_size(new_geom.size)
self.update_color(new_geom.rgba)
self.update_transparency(new_geom.rgba[3])
def update_pos(self, new_pos): # setting the image texture attributes
pos = tuple([float(x) for x in new_pos]) image_shader.CreateIdAttr("UsdUVTexture")
if not self.translate_op: image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(self.texture_file)
self.translate_op = self.xform.AddTranslateOp() image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set("sRGB")
self.translate_op.Set(pos) 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)
def update_rotation(self, new_mat): # setting uvmap shader attributes
r = R.from_matrix(new_mat) uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2")
euler_rotation = r.as_euler('xyz', degrees=True) uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap")
rotation = Gf.Vec3f(float(euler_rotation[0]), float(euler_rotation[1]), float(euler_rotation[2])) uvmap_shader.CreateOutput('results', Sdf.ValueTypeNames.Float2)
if not self.rotate_op: else:
self.rotate_op = self.xform.AddRotateXYZOp() bsdf_shader = UsdShade.Shader.Define(self.stage, mtl_path.AppendPath("Principled_BSDF"))
self.rotate_op.Set(rotation)
# TODO: check to make sure scale and size are the same thing # settings the bsdf shader attributes
def update_size(self, new_size): bsdf_shader.CreateIdAttr("UsdPreviewSurface")
size = tuple([float(x) for x in new_size]) bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(tuple(self.rgba[:3]))
if not self.scale_op: bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(float(self.rgba[-1]))
self.scale_op = self.xform.AddScaleOp() bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
self.scale_op.Set(value=size) bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
def update_color(self, new_color): mtl.CreateSurfaceOutput().ConnectToSource(bsdf_shader.ConnectableAPI(), "surface")
# new_color is the rgba (we extract first three)
rgba = [(float(x) for x in new_color[:3])]
self.prim.GetDisplayColorAttr().Set(rgba)
def update_transparency(self, new_transparency): self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
self.prim.GetDisplayOpacityAttr().Set([float(abs(new_transparency))]) UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
if new_transparency < 0:
self.prim.MakeInvisible()
def __str__(self): def update(
return f'type = {self.type} \ngeom = {self.geom}' self,
pos: np.array,
mat: np.array,
frame: int
):
transformation_mat = create_transform_matrix(rotation_matrix=mat, translation_vector=pos).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
class USDPlane(USDGeom): class USDPrimitiveMesh:
"""
Stores information regarding a plane geom in USD
"""
plane_count = 0
def __init__(self, def __init__(self,
geom=None, stage: Usd.Stage,
stage=None, geom: _structs.MjvGeom,
texture_file=None): objid: int,
super().__init__(geom, stage, texture_file) rgba: List[int] = [1,1,1,1],
self.type = 0 texture_file: Optional[str] = None):
USDPlane.plane_count += 1 self.stage = stage
xform_path = f'/World/Plane_Xform_{USDPlane.plane_count}' self.geom = geom
plane_path = f'{xform_path}/Plane_{USDPlane.plane_count}' self.objid = objid
self.xform = UsdGeom.Xform.Define(stage, xform_path) self.rgba = rgba
self.prim = UsdGeom.Cube.Define(stage, plane_path) # temporary fix for planes
self.ref = stage.GetPrimAtPath(plane_path)
self.set_texture()
def update_size(self, new_size):
size = tuple([float(new_size[0]), float(new_size[1]), 0.005])
if not self.scale_op:
self.scale_op = self.xform.AddScaleOp()
self.scale_op.Set(value=size)
class USDSphere(USDGeom):
"""
Stores information regarding a sphere geom in USD
"""
sphere_count = 0
def __init__(self,
geom=None,
stage=None,
texture_file=None):
super().__init__(geom, stage, texture_file)
self.type = 2
USDSphere.sphere_count += 1
xform_path = f'/World/Sphere_Xform_{USDSphere.sphere_count}'
sphere_path = f'{xform_path}/Sphere_{USDSphere.sphere_count}'
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdGeom.Sphere.Define(stage, sphere_path)
self.ref = stage.GetPrimAtPath(sphere_path)
self.set_texture()
class USDCapsule(USDGeom):
"""
Stores information regarding a capsule geom in USD
"""
capsule_count = 0
def __init__(self,
geom=None,
stage=None,
texture_file=None):
super().__init__(geom, stage, texture_file)
self.type = 3
USDCapsule.capsule_count += 1
xform_path = f'/World/Capsule_Xform_{USDCapsule.capsule_count}'
capsule_path = f'{xform_path}/Capsule_{USDCapsule.capsule_count}'
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdGeom.Capsule.Define(stage, capsule_path)
self.ref = stage.GetPrimAtPath(capsule_path)
self.set_texture()
class USDCylinder(USDGeom):
"""
Stores information regarding a capsule geom in USD
"""
cylinder_count = 0
def __init__(self,
geom=None,
stage=None,
texture_file=None):
super().__init__(geom, stage, texture_file)
self.type = 5
USDCylinder.cylinder_count += 1
xform_path = f'/World/Cylinder_Xform_{USDCylinder.cylinder_count}'
cylinder_path = f'{xform_path}/Cylinder_{USDCylinder.cylinder_count}'
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdGeom.Cylinder.Define(stage, cylinder_path)
self.ref = stage.GetPrimAtPath(cylinder_path)
self.set_texture()
class USDCube(USDGeom):
"""
Stores information regarding a cube geom in USD
"""
cube_count = 0
def __init__(self,
geom=None,
stage=None,
texture_file=None):
super().__init__(geom, stage, texture_file)
self.type = 6
USDCube.cube_count += 1
xform_path = f'/World/Cube_Xform_{USDCube.cube_count}'
cube_path = f'{xform_path}/Cube_{USDCube.cube_count}'
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdGeom.Cube.Define(stage, cube_path)
self.ref = stage.GetPrimAtPath(cube_path)
self.set_texture()
class USDMesh(USDGeom):
"""
Stores information regarding a mesh geom in USD
"""
mesh_count = 0
def __init__(self,
mesh_idx,
geom,
stage,
model,
texture_file):
super().__init__(geom, stage)
assert mesh_idx != -1
mesh_vert_adr_from = model.mesh_vertadr[mesh_idx]
mesh_vert_adr_to = model.mesh_vertadr[mesh_idx+1] if mesh_idx < model.nmesh - 1 else len(model.mesh_vert)
mesh_vert = model.mesh_vert[mesh_vert_adr_from:mesh_vert_adr_to]
mesh_face_adr_from = model.mesh_faceadr[mesh_idx]
mesh_face_adr_to = model.mesh_faceadr[mesh_idx+1] if mesh_idx < model.nmesh - 1 else len(model.mesh_face)
mesh_face = model.mesh_face[mesh_face_adr_from:mesh_face_adr_to]
self.type = 7
xform_path = f'/World/Mesh_Xform_{USDMesh.mesh_count}'
mesh_path= f'{xform_path}/Mesh_{USDMesh.mesh_count}'
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdGeom.Mesh.Define(stage, mesh_path)
self.ref = stage.GetPrimAtPath(mesh_path)
self.vertices = mesh_vert
self.prim.GetPointsAttr().Set(self.vertices)
model.mesh_facenum[mesh_idx]
self.prim.GetFaceVertexCountsAttr().Set([3 for _ in range(model.mesh_facenum[mesh_idx])])
self.faces = mesh_face
self.prim.GetFaceVertexIndicesAttr().Set(self.faces)
self.texture_file = texture_file self.texture_file = texture_file
if texture_file:
mesh_texcoord_adr_from = model.mesh_texcoordadr[mesh_idx] self.prim_mesh = None
mesh_texcoord_adr_to = model.mesh_texcoordadr[mesh_idx+1] if mesh_idx < model.nmesh - 1 else len(model.mesh_texcoord)
mesh_texcoord = model.mesh_texcoord[mesh_texcoord_adr_from:mesh_texcoord_adr_to]
# texid = geom.texid def _get_uv_geometry(self):
# texcoords = model.mesh_texcoord[mesh_texcoord_ranges[texid]:mesh_texcoord_ranges[texid+1]]
mesh_facetexcoord_ranges = get_facetexcoord_ranges(model.nmesh, model.mesh_facenum) assert self.prim_mesh
facetexcoords = model.mesh_facetexcoord.flatten() x_scale, y_scale = self.geom.texrepeat
facetexcoords = facetexcoords[mesh_facetexcoord_ranges[mesh_idx]:mesh_facetexcoord_ranges[mesh_idx+1]]
self.texcoords = UsdGeom.PrimvarsAPI(self.prim).CreatePrimvar("st", mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
mesh_texcoord[:, 0] *= x_scale
mesh_texcoord[:, 1] *= y_scale
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:
# remove all code in this if block
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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
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.array,
mat: np.array,
frame: int
):
transformation_mat = create_transform_matrix(rotation_matrix=mat, translation_vector=pos).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
class USDPrimitive:
def __init__(self,
stage: Usd.Stage,
geom: _structs.MjvGeom,
objid: int,
rgba: List[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
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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
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.array,
mat: np.array,
frame: int
):
transformation_mat = create_transform_matrix(rotation_matrix=mat, translation_vector=pos).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
class USDCapsule(USDPrimitive):
def __init__(self,
stage: Usd.Stage,
geom: _structs.MjvGeom,
objid: int,
rgba: List[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()
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: _structs.MjvGeom,
objid: int,
rgba: List[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()
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: _structs.MjvGeom,
objid: int,
rgba: List[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, Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying) UsdGeom.Tokens.faceVarying)
self.texcoords.Set(mesh_texcoord) self.texcoords.Set(mesh_texcoord)
self.texcoords.SetIndices(Vt.IntArray(facetexcoords.tolist())); self.texcoords.SetIndices(Vt.IntArray([i for i in range(mesh_facenum*3)]))
mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(texture_file))[0]}") # setting attributes for the shape
mtl = UsdShade.Material.Define(stage, mtl_path) self._attach_material()
shader = UsdShade.Shader.Define(stage, mtl_path.AppendPath("Shader"))
shader.CreateIdAttr("UsdPreviewSurface")
shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0))
shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
diffuse_tx = UsdShade.Shader.Define(stage,mtl_path.AppendPath("DiffuseColorTx")) # defining ops required by update function
diffuse_tx.CreateIdAttr('UsdUVTexture') self.transform_op = self.usd_xform.AddTransformOp()
diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(texture_file)
diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3)
shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb')
mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface")
self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.prim).Bind(mtl)
USDMesh.mesh_count += 1
def update_geom(self, new_geom):
self.update_pos(new_geom.pos)
self.update_rotation(new_geom.mat)
if not self.texture_file:
self.update_color(new_geom.rgba)
self.update_transparency(new_geom.rgba[3])
class USDLight(object):
"""
Class for the created lights
"""
light_count = 0
class USDSphereMesh(USDPrimitiveMesh):
def __init__(self, def __init__(self,
stage): stage: Usd.Stage,
self.stage = stage geom: _structs.MjvGeom,
USDLight.light_count += 1 objid: int,
xform_path = f'/World/Light_Xform_{USDLight.light_count}' rgba: List[int] = [1,1,1,1],
light_path = f'{xform_path}/Light_{USDLight.light_count}' texture_file: Optional[str] = None):
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdLux.SphereLight.Define(stage, light_path)
self.ref = stage.GetPrimAtPath(light_path)
self.translate_op = None super().__init__(stage,
geom,
objid,
rgba,
texture_file)
def update_light(self, new_light): xform_path = f'/World/SphereMesh_Xform_{objid}'
pos = tuple([float(x) for x in new_light.pos]) mesh_path= f'{xform_path}/SphereMesh_{objid}'
if not self.translate_op: self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.translate_op = self.xform.AddTranslateOp() self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.translate_op.Set(pos) self.usd_prim = stage.GetPrimAtPath(mesh_path)
if pos == (0, 0, 0): self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere(radius=float(self.geom.size[0]),
self.prim.GetIntensityAttr().Set(0); create_uv_map=True)
else:
self.prim.GetIntensityAttr().Set(20000);
# TODO attributes: self.prim_mesh.translate(-self.prim_mesh.get_center())
# - direction
# - intensity
# - exposure
# - radius
# - specular
class USDCamera(object): mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
""" self.usd_mesh.GetPointsAttr().Set(mesh_vert)
Class for created cameras self.usd_mesh.GetFaceVertexCountsAttr().Set([3 for _ in range(mesh_facenum)])
""" self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
camera_count = 0 # 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)]))
# 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, def __init__(self,
stage): stage: Usd.Stage,
geom: _structs.MjvGeom,
objid: int,
rgba: List[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)]))
# 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: _structs.MjvGeom,
objid: int,
rgba: List[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)]))
# setting attributes for the shape
self._attach_material()
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
class USDLight:
def __init__(self,
stage: Usd.Stage,
objid: int,
radius: Optional[float] = 0.7):
self.stage = stage self.stage = stage
USDCamera.camera_count += 1
camera_path = f'/World/Camera_{USDCamera.camera_count}'
self.prim = UsdGeom.Camera.Define(stage, camera_path)
self.ref = stage.GetPrimAtPath(camera_path)
def update_camera(self, new_pos, new_quat): xform_path = f'/World/Light_Xform_{objid}'
# print("---- Updating camera in USD ----") light_path = f'{xform_path}/Light_{objid}'
xformAPI = UsdGeom.XformCommonAPI(self.prim) 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)
pos = tuple([float(x) for x in new_pos]) # we assume in mujoco that all lights are point lights
self.usd_light.GetRadiusAttr().Set(radius)
# convert a quat to euler rotation angles self.usd_light.GetTreatAsPointAttr().Set(False)
r = R.from_quat(new_quat) self.usd_light.GetNormalizeAttr().Set(True)
euler_rotation = r.as_euler('xyz', degrees=True)
rotation = Gf.Vec3f(float(euler_rotation[2]), float(euler_rotation[1]), float(euler_rotation[0]))
# hardcoded values for Robosuite testing and prototype
# TODO: use actual camera values
xformAPI.SetTranslate(pos)
xformAPI.SetRotate(rotation)
# xformAPI.SetScale((1, 1, 1))
self.prim.CreateFocalLengthAttr().Set(24)
self.prim.CreateFocusDistanceAttr().Set(400)
# defining ops required by update function
self.translate_op = self.usd_xform.AddTranslateOp()
def update(self,
pos: np.array,
intensity: int,
color: np.array,
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 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.CreateFocusDistanceAttr().Set(400)
self.usd_camera.GetClippingRangeAttr().Set(Gf.Vec2f(1e-4, 1e6))
def update(self,
cam_pos: np.array,
cam_mat: np.array,
frame: int):
transformation_mat = create_transform_matrix(rotation_matrix=cam_mat, translation_vector=cam_pos).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
python/mujoco/usd_renderer.py
+282 -174
View File
@@ -1,47 +1,87 @@
import os import os
import shutil import pprint
from termcolor import colored
import mujoco import mujoco
import mujoco.viewer as viewer from usd_utils import *
from mujoco.usd_component import *
from mujoco.usd_utilities import *
from pxr import Usd, UsdGeom
from mujoco import _structs
from PIL import Image as im
from PIL import ImageOps from PIL import ImageOps
from mujoco import mjtGeom
from PIL import Image as im
from usd_component import *
from pxr import Usd, UsdGeom
from termcolor import colored
from mujoco import mjv_averageCamera
from typing import Optional, List, Union, Tuple
from mujoco import _structs, _constants, _enums
from scipy.spatial.transform import Rotation as R
class USDRenderer(object): class USDRenderer:
def __init__(
self,
model: _structs.MjModel,
height: int = 480,
width: int = 480,
max_geom: int = 10000,
output_directory_name: str = "mujoco_usdpkg",
output_directory_root: str = "./",
verbose: bool = True,
light_intensity: int = 10000
):
""" Initializes a new USD Renderer
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.
""" """
Renderer class that creates USD representations for mujoco scenes
""" buffer_width = model.vis.global_.offwidth
def __init__(self, buffer_height = model.vis.global_.offheight
model,
height=480, if width > buffer_width:
width=480, raise ValueError(f"""
root_dir_name="usdpkg", Image width {width} > framebuffer width {buffer_width}. Either reduce the image
root_dir_path=None, width or specify a larger offscreen framebuffer in the model XML using the
verbose=True): 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.model = model
self.root_dir_name = root_dir_name self.height = height
self.root_dir_path = root_dir_path self.width = width
self.max_geom = max_geom
self.output_directory_name = output_directory_name
self.output_directory_root = output_directory_root
self.verbose = verbose self.verbose = verbose
self.data = None self.light_intensity = light_intensity
self.renderer = mujoco.Renderer(model, height, width)
self.reload_scene_info = True
self.frame_count = 0
self.create_output_directories() self.frame_count = 0 # maintains how many times we have saved the scene
self.updates = 0
self.stage = Usd.Stage.CreateInMemory() # initializing rendering requirements
self.renderer = mujoco.Renderer(model, height, width, max_geom)
self._initialize_usd_stage()
self._scene_option = _structs.MjvOption() # using default scene option
UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z) # initializing output_directories
self._initialize_output_directories()
geom_groups = [0,1,0,0,0,0] # Setting default geom groups for now # loading required textures for the scene
self._load_textures()
self.scene_option = _structs.MjvOption()
self.scene_option.geomgroup = geom_groups
@property @property
def usd(self): def usd(self):
@@ -51,167 +91,235 @@ class USDRenderer(object):
def scene(self): def scene(self):
return self.renderer.scene return self.renderer.scene
def create_output_directories(self): def _initialize_usd_stage(self):
if not self.root_dir_path: self.stage = Usd.Stage.CreateInMemory()
self.root_dir_path = os.getcwd() UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
self.stage.SetStartTimeCode(0)
# add as user imput
self.stage.SetTimeCodesPerSecond(60.0)
self.output_dir = os.path.join(self.root_dir_path, self.root_dir_name) def _initialize_output_directories(self):
if not os.path.exists(self.output_dir): self.output_directory_path = os.path.join(self.output_directory_root, self.output_directory_name)
os.makedirs(self.output_dir) if not os.path.exists(self.output_directory_path):
os.makedirs(self.output_directory_path)
self.scenes_dir = os.path.join(self.output_dir, "scenes") self.frames_directory = os.path.join(self.output_directory_path, "frames")
if not os.path.exists(self.scenes_dir): if not os.path.exists(self.frames_directory):
os.makedirs(self.scenes_dir) os.makedirs(self.frames_directory)
self.assets_dir = os.path.join(self.output_dir, "assets") self.assets_directory = os.path.join(self.output_directory_path, "assets")
if not os.path.exists(self.assets_dir): if not os.path.exists(self.assets_directory):
os.makedirs(self.assets_dir) os.makedirs(self.assets_directory)
if self.verbose: if self.verbose:
output_dir_msg = colored(f"Writing files to {self.output_dir}", "green") print(colored(f"Writing output frames and assets to {self.output_directory_path}", "green"))
print(output_dir_msg)
def update_scene(
self,
data: _structs.MjData,
camera: Union[int, str, _structs.MjvCamera] = -1,
scene_option: Optional[_structs.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
"""
def save_scene(self):
output_file_path = os.path.join(self.scenes_dir, f'frame_{self.frame_count}_.usd')
with open(output_file_path, "w") as f:
f.write(self.usd)
self.frame_count += 1 self.frame_count += 1
def update_geom_groups(self, geom_groups): scene_option = scene_option or self._scene_option
self.scene_option.geomgroup = geom_groups
self.reload_scene_info = True
self.update_scene(self.data)
def update_scene(self, data): # update the mujoco renderer
self.renderer.update_scene(data, scene_option=self.scene_option) self.renderer.update_scene(data,
self.data = data scene_option=scene_option,
camera=camera)
if self.reload_scene_info: # TODO: update scene options
# loads the initial geoms, lights, and camera information if self.updates == 0:
# from the scene self._initialize_usd_stage()
self._load()
self.reload_scene_info = False
self._update() self._load_geoms()
self._load_lights()
self._load_cameras()
def _load(self):
"""
Loads and initializes the necessary objects to render the scene
"""
# Create and loads the texture files to the assets directory
# TODO: remove code once added internally to mujoco
data_adr = 0
texture_files = []
for texid in range(self.model.ntex):
height = self.model.tex_height[texid]
width = self.model.tex_width[texid]
pixels = 3*height*width
rgb = self.model.tex_rgb[data_adr:data_adr+pixels]
img = rgb.reshape(height, width, 3)
texture_file_name = f"texture_{texid}.png"
file_path = os.path.join(self.assets_dir, texture_file_name)
img = im.fromarray(img)
img = ImageOps.flip(img)
img.save(file_path)
relative_path = os.path.relpath(self.assets_dir, self.scenes_dir)
img_path = os.path.join(relative_path, texture_file_name)
texture_files.append(img_path)
data_adr += pixels
# initializes an array to store all the geoms in the scene
# populates with "empty" USDGeom objects
self.usd_geoms = []
geoms = self.scene.geoms
self.ngeom = self.scene.ngeom
for i in range(self.ngeom):
geom = geoms[i]
if geom.texid == -1:
texture_file = None
else:
texture_file = texture_files[geom.texid]
if geom.type == USDGeomType.Mesh.value:
self.usd_geoms.append(USDMesh(self.model.geom_dataid[geom.objid],
geom,
self.stage,
self.model,
texture_file))
else:
self.usd_geoms.append(create_usd_geom_primitive(geom,
self.stage,
texture_file))
# initializes an array to store all the lights in the scene
# populates with "empty" USDLight objects
self.usd_lights = []
lights = self.scene.lights
self.nlight = self.scene.nlight
for i in range(self.nlight):
self.usd_lights.append(USDLight(self.stage))
# initializes an array to store all the cameras in the scene
# populates with "empty" USDCamera objects
self.usd_cameras = []
ncam = self.model.ncam
for i in range(ncam):
self.usd_cameras.append(USDCamera(self.stage))
def _update(self):
self._update_geoms() self._update_geoms()
self._update_lights() self._update_lights()
self._update_cameras() self._update_cameras()
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 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
# absolute path for cluster, TODO: remove!
abs_path = os.path.join(os.path.abspath(self.assets_directory), texture_file_name)
self.texture_files.append(abs_path)
data_adr += pixels
if self.verbose:
print(colored(f"Writing texture {texture_id}", "cyan"))
if self.verbose:
print(colored(f"Completed writing {self.model.ntex} textures to {self.assets_directory}", "green"))
def _load_geoms(self):
# stores a list of all the geoms in the scene
self.usd_geoms = []
# initializing the geoms
for i in range(self.scene.ngeom):
geom = self.scene.geoms[i]
if geom.rgba[3] <= 0:
self.usd_geoms.append(None)
continue
# handles meshes in scene
if geom.type == mjtGeom.mjGEOM_MESH:
usd_geom = USDMesh(stage=self.stage,
model=self.model,
geom=geom,
objid=i,
dataid=self.model.geom_dataid[geom.objid],
rgba=geom.rgba,
texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
# handles primitives
else:
if geom.type == mjtGeom.mjGEOM_PLANE:
usd_geom = USDPlaneMesh(stage=self.stage,
geom=geom,
objid=i,
rgba=geom.rgba,
texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
elif geom.type == mjtGeom.mjGEOM_SPHERE:
usd_geom = USDSphereMesh(stage=self.stage,
geom=geom,
objid=i,
rgba=geom.rgba,
texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
elif geom.type == mjtGeom.mjGEOM_CAPSULE:
usd_geom = USDCapsule(stage=self.stage,
geom=geom,
objid=i,
rgba=geom.rgba,
texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
elif geom.type == mjtGeom.mjGEOM_ELLIPSOID:
usd_geom = USDEllipsoid(stage=self.stage,
geom=geom,
objid=i,
rgba=geom.rgba,
texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
elif geom.type == mjtGeom.mjGEOM_CYLINDER:
usd_geom = USDCylinderMesh(stage=self.stage,
geom=geom,
objid=i,
rgba=geom.rgba,
texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
elif geom.type == mjtGeom.mjGEOM_BOX:
usd_geom = USDCubeMesh(stage=self.stage,
geom=geom,
objid=i,
rgba=geom.rgba,
texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
else:
usd_geom = None
self.usd_geoms.append(usd_geom)
def _update_geoms(self): def _update_geoms(self):
""" # iterate through all geoms in the scene and makes update
Updates the geoms to match the current scene for i in range(self.scene.ngeom):
""" geom = self.scene.geoms[i]
geoms = self.scene.geoms if self.usd_geoms[i]:
for i in range(self.ngeom): self.usd_geoms[i].update(pos=geom.pos,
if self.usd_geoms[i]: # TODO: remove this once all primitives are added mat=geom.mat,
self.usd_geoms[i].update_geom(geoms[i]) frame=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]
self.usd_lights.append(USDLight(stage=self.stage,
objid=i))
def _update_lights(self): def _update_lights(self):
""" for i in range(self.scene.nlight):
Updates the lights to match the current scene light = self.scene.lights[i]
""" self.usd_lights[i].update(pos=light.pos,
lights = self.scene.lights intensity=self.light_intensity,
nlight = self.scene.nlight color=light.diffuse,
for i in range(nlight): frame=self.updates)
self.usd_lights[i].update_light(lights[i])
def _load_cameras(self):
self.camera = USDCamera(stage=self.stage,
objid=0)
def _update_cameras(self): def _update_cameras(self):
""" camera = mjv_averageCamera(self.scene.camera[0], self.scene.camera[1])
Updates the camera to match the current scene
"""
ncam = self.model.ncam
for i in range(ncam):
self.usd_cameras[i].update_camera(self.model.cam_pos[i], self.model.cam_quat[i])
def compress(self): forward = camera.forward
""" up = camera.up
Compresses the output directory to a zip file for easy transfer right = np.cross(forward, up)
"""
R = np.eye(3)
R[:, 0] = right
R[:, 1] = up
R[:, 2] = -forward
self.camera.update(cam_pos=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.array] = np.array([0.3, 0.3, 0.3]),
objid: Optional[int]=1):
new_light = USDLight(stage=self.stage,
objid=objid,
radius=radius)
new_light.update(pos=pos,
intensity=intensity,
color=color,
frame=0)
def add_camera(self,
pos:List[float],
rotation_xyz:List[float],
objid: Optional[int]=1):
# TODO: change this!
new_camera = USDCamera(stage=self.stage,
objid=objid)
r = R.from_euler('xyz', rotation_xyz, degrees=True)
new_camera.update(cam_pos=pos,
cam_mat=r.as_matrix(),
frame=0)
def save_scene(self):
self.stage.SetEndTimeCode(self.frame_count)
# with open(f'./{self.output_directory_name}/frames/frame_{self.frame_count}_.usd', "w") as f:
# f.write(self.usd)
self.stage.Export(f'./{self.output_directory_name}/frames/frame_{self.frame_count}_.usd')
if self.verbose: if self.verbose:
output_dir_msg = colored(f"Compressing files at {self.output_dir} and saving at {self.output_dir}", "green") print(colored(f"Writing frame_{self.frame_count}", "green"))#
print(output_dir_msg)
shutil.make_archive(base_name=self.output_dir,
format='zip',
base_dir=self.root_dir_name)
def start_viewer(self):
if self.data:
viewer.launch(self.model)
def render(self):
# should render the usd file given a particular renderer that
# works with USD files
# TODO: determine if this is valid functionality
pass
# TODO: remove later, this is only for debugging purposes
def print_geom_information(self):
for i in range(self.ngeom):
print(self.usd_geoms[i])
python/mujoco/usd_utilities.py
-16
View File
@@ -1,16 +0,0 @@
def get_mesh_ranges(nmesh, arr):
mesh_ranges = [0]
running_sum = 0
for i in range(nmesh):
running_sum += arr[i]
mesh_ranges.append(running_sum)
return mesh_ranges
def get_facetexcoord_ranges(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
python/mujoco/usd_utils.py
+12
View File
@@ -0,0 +1,12 @@
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