diff --git a/python/dist/mujoco-3.0.1.tar.gz b/python/dist/mujoco-3.0.1.tar.gz
deleted file mode 100644
index ce307b11..00000000
Binary files a/python/dist/mujoco-3.0.1.tar.gz and /dev/null differ
diff --git a/python/dist/mujoco-3.0.2.tar.gz b/python/dist/mujoco-3.0.2.tar.gz
deleted file mode 100644
index cd71ac7c..00000000
Binary files a/python/dist/mujoco-3.0.2.tar.gz and /dev/null differ
diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py
index d9e0e5cf..bad3dad3 100644
--- a/python/mujoco/__init__.py
+++ b/python/mujoco/__init__.py
@@ -46,6 +46,9 @@ from mujoco._render import *
from mujoco._structs import *
from mujoco.gl_context import *
from mujoco.renderer import Renderer
+from mujoco.usd.usd_exporter import USDExporter
+from mujoco.usd.usd_component import *
+from mujoco.usd.usd_utils import *
HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')
diff --git a/python/mujoco/usd_component.py b/python/mujoco/usd/usd_component.py
similarity index 74%
rename from python/mujoco/usd_component.py
rename to python/mujoco/usd/usd_component.py
index c649096e..7c7e74a8 100644
--- a/python/mujoco/usd_component.py
+++ b/python/mujoco/usd/usd_component.py
@@ -1,12 +1,15 @@
+from typing import Optional, List, Tuple
+
import numpy as np
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
+from mujoco import mjtGeom
+from mujoco.usd.usd_utils import *
+from mujoco import _structs, _enums
+
class USDMesh:
def __init__(
@@ -16,7 +19,7 @@ class USDMesh:
geom: _structs.MjvGeom,
objid: int,
dataid: int,
- rgba: List[int] = [1,1,1,1],
+ rgba: Tuple[int] = (1,1,1,1),
texture_file: Optional[str] = None
):
""" Initializes a new USD mesh
@@ -58,7 +61,11 @@ class USDMesh:
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
- def get_facetexcoord_ranges(self, nmesh, arr):
+ def get_facetexcoord_ranges(
+ self,
+ nmesh,
+ arr
+ ):
facetexcoords_ranges = [0]
running_sum = 0
for i in range(nmesh):
@@ -68,24 +75,31 @@ class USDMesh:
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_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):
- # 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_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_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]
@@ -105,8 +119,8 @@ class USDMesh:
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)
+ 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")
@@ -131,10 +145,11 @@ class USDMesh:
# settings the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
- bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(tuple(self.rgba[:3]))
+
+ bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(tuple(color))
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)
+ 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")
@@ -145,19 +160,33 @@ class USDMesh:
self,
pos: np.array,
mat: np.array,
+ visible: bool,
frame: int
):
transformation_mat = 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: _structs.MjvGeom,
- objid: int,
- rgba: List[int] = [1,1,1,1],
- texture_file: Optional[str] = None):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: Tuple[int] = (1, 1, 1, 1),
+ texture_file: Optional[str] = None
+ ):
self.stage = stage
self.geom = geom
self.objid = objid
@@ -166,6 +195,9 @@ class USDPrimitiveMesh:
self.prim_mesh = None
+ def _set_refinement_properties(self):
+ self.usd_prim.GetAttribute('subdivisionScheme').Set("none")
+
def _get_uv_geometry(self):
assert self.prim_mesh
@@ -175,8 +207,12 @@ class USDPrimitiveMesh:
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
+ 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()
@@ -194,7 +230,6 @@ class USDPrimitiveMesh:
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"))
@@ -203,8 +238,8 @@ class USDPrimitiveMesh:
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)
+ 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")
@@ -231,8 +266,8 @@ class USDPrimitiveMesh:
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)
+ 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")
@@ -241,25 +276,42 @@ class USDPrimitiveMesh:
def update(
self,
- pos: np.array,
- mat: np.array,
+ pos: np.ndarray,
+ mat: np.ndarray,
+ visible: bool,
frame: int
):
transformation_mat = 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: _structs.MjvGeom,
- objid: int,
- rgba: List[int] = [1,1,1,1],
- texture_file: Optional[str] = None):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.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
+
+ 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}")
@@ -273,8 +325,8 @@ class USDPrimitive:
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)
+ 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")
@@ -299,8 +351,8 @@ class USDPrimitive:
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)
+ 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")
@@ -309,26 +361,42 @@ class USDPrimitive:
def update(
self,
- pos: np.array,
- mat: np.array,
+ pos: np.ndarray,
+ mat: np.ndarray,
+ visible: bool,
frame: int
):
transformation_mat = 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: _structs.MjvGeom,
- objid: int,
- rgba: List[int] = [1,1,1,1],
- texture_file: Optional[str] = None):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: Tuple[int] = (1, 1, 1, 1),
+ texture_file: Optional[str] = None
+ ):
- super().__init__(stage,
- geom,
- objid,
- rgba,
- texture_file)
+ super().__init__(
+ stage,
+ geom,
+ objid,
+ rgba,
+ texture_file
+ )
xform_path = f'/World/Capsule_Xform_{objid}'
capsule_path = f'{xform_path}/Capsule_{objid}'
@@ -344,23 +412,29 @@ class USDCapsule(USDPrimitive):
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: _structs.MjvGeom,
- objid: int,
- rgba: List[int] = [1,1,1,1],
- texture_file: Optional[str] = None):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: Tuple[int] = (1, 1, 1, 1),
+ texture_file: Optional[str] = None
+ ):
- super().__init__(stage,
- geom,
- objid,
- rgba,
- texture_file)
+ super().__init__(
+ stage,
+ geom,
+ objid,
+ rgba,
+ texture_file
+ )
xform_path = f'/World/Ellipsoid_Xform_{objid}'
ellipsoid_path = f'{xform_path}/Ellipsoid_{objid}'
@@ -376,22 +450,28 @@ class USDEllipsoid(USDPrimitive):
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: _structs.MjvGeom,
- objid: int,
- rgba: List[int] = [1,1,1,1],
- texture_file: Optional[str] = None):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: Tuple[int] = (1, 1, 1, 1),
+ texture_file: Optional[str] = None
+ ):
- super().__init__(stage,
- geom,
- objid,
- rgba,
- texture_file)
+ super().__init__(
+ stage,
+ geom,
+ objid,
+ rgba,
+ texture_file
+ )
xform_path = f'/World/CubeMesh_Xform_{objid}'
mesh_path= f'{xform_path}/CubeMesh_{objid}'
@@ -399,11 +479,13 @@ class USDCubeMesh(USDPrimitiveMesh):
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 = 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())
@@ -421,6 +503,8 @@ class USDCubeMesh(USDPrimitiveMesh):
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()
@@ -428,18 +512,22 @@ class USDCubeMesh(USDPrimitiveMesh):
self.transform_op = self.usd_xform.AddTransformOp()
class USDSphereMesh(USDPrimitiveMesh):
- def __init__(self,
- stage: Usd.Stage,
- geom: _structs.MjvGeom,
- objid: int,
- rgba: List[int] = [1,1,1,1],
- texture_file: Optional[str] = None):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: Tuple[int] = (1, 1, 1, 1),
+ texture_file: Optional[str] = None
+ ):
- super().__init__(stage,
- geom,
- objid,
- rgba,
- texture_file)
+ super().__init__(
+ stage,
+ geom,
+ objid,
+ rgba,
+ texture_file
+ )
xform_path = f'/World/SphereMesh_Xform_{objid}'
mesh_path= f'{xform_path}/SphereMesh_{objid}'
@@ -447,8 +535,10 @@ class USDSphereMesh(USDPrimitiveMesh):
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 = o3d.geometry.TriangleMesh.create_sphere(
+ radius=float(self.geom.size[0]),
+ create_uv_map=True
+ )
self.prim_mesh.translate(-self.prim_mesh.get_center())
@@ -466,6 +556,8 @@ class USDSphereMesh(USDPrimitiveMesh):
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()
@@ -473,18 +565,22 @@ class USDSphereMesh(USDPrimitiveMesh):
self.transform_op = self.usd_xform.AddTransformOp()
class USDCylinderMesh(USDPrimitiveMesh):
- def __init__(self,
- stage: Usd.Stage,
- geom: _structs.MjvGeom,
- objid: int,
- rgba: List[int] = [1,1,1,1],
- texture_file: Optional[str] = None):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: Tuple[int] = (1, 1, 1, 1),
+ texture_file: Optional[str] = None
+ ):
- super().__init__(stage,
- geom,
- objid,
- rgba,
- texture_file)
+ super().__init__(
+ stage,
+ geom,
+ objid,
+ rgba,
+ texture_file
+ )
xform_path = f'/World/CylinderMesh_Xform_{objid}'
mesh_path= f'{xform_path}/CylinderMesh_{objid}'
@@ -492,9 +588,11 @@ class USDCylinderMesh(USDPrimitiveMesh):
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 = 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())
@@ -512,6 +610,8 @@ class USDCylinderMesh(USDPrimitiveMesh):
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()
@@ -519,18 +619,22 @@ class USDCylinderMesh(USDPrimitiveMesh):
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)
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ geom: _structs.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}'
@@ -538,11 +642,13 @@ class USDPlaneMesh(USDPrimitiveMesh):
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 = 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())
@@ -560,17 +666,21 @@ class USDPlaneMesh(USDPrimitiveMesh):
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 USDLight:
- def __init__(self,
- stage: Usd.Stage,
- objid: int,
- radius: Optional[float] = 0.7):
+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}'
@@ -587,11 +697,13 @@ class USDLight:
# 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):
+ 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):
@@ -600,10 +712,40 @@ class USDLight:
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):
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ objid: int
+ ):
self.stage = stage
xform_path = f'/World/Camera_Xform_{objid}'
@@ -615,15 +757,20 @@ class USDCamera:
# 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(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.array,
- cam_mat: np.array,
- frame: int):
+ def update(
+ self,
+ cam_pos: np.ndarray,
+ cam_mat: np.ndarray,
+ 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)
diff --git a/python/mujoco/usd/usd_exporter.py b/python/mujoco/usd/usd_exporter.py
new file mode 100644
index 00000000..020c8d34
--- /dev/null
+++ b/python/mujoco/usd/usd_exporter.py
@@ -0,0 +1,385 @@
+import os
+import json
+from typing import Optional, List, Union, Tuple
+
+import pprint
+from tqdm import tqdm
+from PIL import ImageOps
+from PIL import Image as im
+from pxr import Gf, Sdf, Vt
+from pxr import Usd, UsdGeom
+from termcolor import colored
+from scipy.spatial.transform import Rotation as R
+
+import mujoco
+from mujoco import mjtGeom
+from mujoco.usd.usd_utils import *
+from mujoco.usd.usd_component import *
+from mujoco import mjv_averageCamera
+from mujoco import _structs, _constants, _enums
+
+
+class USDExporter:
+
+ 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 = "./",
+ light_intensity: int = 10000,
+ camera_names: List[str] = None,
+ specialized_materials_file: str = None,
+ verbose: bool = True,
+ ):
+ """ 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.
+ """
+
+ 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:
+
+
+ """.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:
+
+
+ """.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
+
+ # assert specialized_materials_file.endswith('.json')
+ # self.specialized_materials = json.loads(specialized_materials_file)
+
+ 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 = _structs.MjvOption() # using default scene option
+
+ # initializing output_directories
+ self._initialize_output_directories()
+
+ # loading required textures for the scene
+ self._load_textures()
+
+ self.extra_added_water = False
+ self.extra_added_stove = False
+ self.extra_added_coffee = False
+
+ @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 imput
+ 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(colored(f"Writing output frames and assets to {self.output_directory_path}", "green"))
+
+ def update_scene(
+ self,
+ data: _structs.MjData,
+ 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
+ """
+
+ 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(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(colored(f"Completed writing {self.model.ntex} textures to {self.assets_directory}", "green"))
+
+ def _load_geom(
+ self,
+ geom: _structs.MjvGeom
+ ):
+
+ geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
+ assert geom_name not in self.geom_name2usd
+
+ # handles meshes in scene
+ if geom.type == mjtGeom.mjGEOM_MESH:
+ usd_geom = USDMesh(stage=self.stage,
+ model=self.model,
+ geom=geom,
+ objid=geom_name,
+ dataid=self.model.geom_dataid[geom.objid],
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ elif geom.type == mjtGeom.mjGEOM_PLANE:
+ usd_geom = USDPlaneMesh(stage=self.stage,
+ geom=geom,
+ objid=geom_name,
+ 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=geom_name,
+ 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=geom_name,
+ 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=geom_name,
+ 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=geom_name,
+ 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=geom_name,
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ 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(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 not np.allclose(light.pos, [0, 0, 0]):
+ 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(USDCamera(stage=self.stage,
+ objid=name))
+
+ def _update_cameras(
+ self,
+ data: _structs.MjData,
+ scene_option: Optional[_structs.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 = 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.array] = np.array([0.3, 0.3, 0.3]),
+ objid: Optional[int]=1,
+ light_type: Optional[str]="sphere"
+ ):
+
+ if light_type == "sphere":
+ new_light = 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 = 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 = 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,
+ 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(colored(f"Writing frame_{self.frame_count}", "green"))
\ No newline at end of file
diff --git a/python/mujoco/usd_utils.py b/python/mujoco/usd/usd_utils.py
similarity index 100%
rename from python/mujoco/usd_utils.py
rename to python/mujoco/usd/usd_utils.py
diff --git a/python/mujoco/usd_exporter.py b/python/mujoco/usd_exporter.py
deleted file mode 100644
index 1a5e3ca8..00000000
--- a/python/mujoco/usd_exporter.py
+++ /dev/null
@@ -1,322 +0,0 @@
-import os
-import pprint
-import mujoco
-from usd_utils import *
-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 USDExporter:
-
- 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.
- """
-
- 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:
-
-
- """.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:
-
-
- """.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.verbose = verbose
- self.light_intensity = light_intensity
-
- self.frame_count = 0 # maintains how many times we have saved the scene
- self.updates = 0
-
- # 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
-
- # 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 imput
- self.stage.SetTimeCodesPerSecond(60.0)
-
- 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(colored(f"Writing output frames and assets to {self.output_directory_path}", "green"))
-
- 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
- """
-
- 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,
- camera=camera)
-
- # TODO: update scene options
- if self.updates == 0:
- self._initialize_usd_stage()
-
- self._load_geoms()
- self._load_lights()
- self._load_cameras()
-
- self._update_geoms()
- self._update_lights()
- 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
-
- self.texture_files.append(img_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):
- # iterate through all geoms in the scene and makes update
- for i in range(self.scene.ngeom):
- geom = self.scene.geoms[i]
- if self.usd_geoms[i]:
- self.usd_geoms[i].update(pos=geom.pos,
- mat=geom.mat,
- 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):
- for i in range(self.scene.nlight):
- light = self.scene.lights[i]
- self.usd_lights[i].update(pos=light.pos,
- intensity=self.light_intensity,
- color=light.diffuse,
- frame=self.updates)
-
- def _load_cameras(self):
- self.camera = USDCamera(stage=self.stage,
- objid=0)
-
- def _update_cameras(self):
- camera = mjv_averageCamera(self.scene.camera[0], self.scene.camera[1])
-
- forward = camera.forward
- up = camera.up
- 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:
- print(colored(f"Writing frame_{self.frame_count}", "green"))#
\ No newline at end of file