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

addressing PR comments

Browse Source
This commit is contained in:
Abhishek Joshi
2024-02-21 10:18:13 -06:00
parent f721d3900d
commit 7ffd43b55c
7 changed files with 677 additions and 464 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dist/mujoco-3.0.1.tar.gz
Vendored
BIN
View File
Binary file not shown.
python/dist/mujoco-3.0.2.tar.gz
Vendored
BIN
View File
Binary file not shown.
python/mujoco/__init__.py
+3
View File
@@ -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')
python/mujoco/usd_component.py → python/mujoco/usd/usd_component.py
+289 -142
View File
@@ -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)
python/mujoco/usd/usd_exporter.py
+385
View File
@@ -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:
<visual>
<global offwidth="my_width"/>
</visual>""".lstrip())
if height > buffer_height:
raise ValueError(f"""
Image height {height} > framebuffer height {buffer_height}. Either reduce the
image height or specify a larger offscreen framebuffer in the model XML using
the clause:
<visual>
<global offheight="my_height"/>
</visual>""".lstrip())
self.model = model
self.height = height
self.width = width
self.max_geom = max_geom
self.output_directory_name = output_directory_name
self.output_directory_root = output_directory_root
self.light_intensity = light_intensity
self.camera_names = camera_names
self.specialized_materials_file = specialized_materials_file
self.verbose = verbose
# 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"))
python/mujoco/usd_utils.py → python/mujoco/usd/usd_utils.py
View File
python/mujoco/usd_exporter.py
-322
View File
@@ -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:
<visual>
<global offwidth="my_width"/>
</visual>""".lstrip())
if height > buffer_height:
raise ValueError(f"""
Image height {height} > framebuffer height {buffer_height}. Either reduce the
image height or specify a larger offscreen framebuffer in the model XML using
the clause:
<visual>
<global offheight="my_height"/>
</visual>""".lstrip())
self.model = model
self.height = height
self.width = width
self.max_geom = max_geom
self.output_directory_name = output_directory_name
self.output_directory_root = output_directory_root
self.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"))#