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

Adding logic to include meshes using MjModel

Browse Source
This commit is contained in:
Abhishek Joshi
2023-09-09 16:37:45 -05:00
parent 4e58b0619b
commit 0412f98df1
4 changed files with 131 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dist/mujoco-2.3.8.tar.gz
Vendored
BIN
View File
Binary file not shown.
python/mujoco/usd_component.py
+98 -31
View File
@@ -1,7 +1,17 @@
from enum import Enum
import mujoco
from pxr import Usd, UsdGeom, Vt
from pxr import Usd, UsdGeom, UsdLux, Vt, Gf
from scipy.spatial.transform import Rotation as R
def create_usd_geom_primitive(geom, stage):
geom_type = geom.type
if geom_type==USDGeomType.Plane.value:
return USDPlane(geom, stage)
elif geom_type==USDGeomType.Sphere.value:
return USDSphere(geom, stage)
elif geom_type==USDGeomType.Cube.value:
return USDCube(geom, stage)
class USDGeomType(Enum):
"""
@@ -12,6 +22,7 @@ class USDGeomType(Enum):
Plane = 0
Sphere = 2
Cube = 6
Mesh = 7
class USDGeom(object):
"""
@@ -28,12 +39,22 @@ class USDGeom(object):
self.ref = None
def update_geom(self, new_geom):
raise NotImplementedError
self.update_pos(new_geom.pos)
self.update_rotation(new_geom.mat)
self.update_size(new_geom.size)
self.update_color(new_geom.rgba)
def update_pos(self, new_pos):
pos = tuple([float(x) for x in new_pos])
self.xform.AddTranslateOp().Set(pos)
def update_rotation(self, new_mat):
r = R.from_matrix(new_mat)
euler_rotation = r.as_euler('xyz', degrees=True)
rotation = Gf.Vec3f(float(euler_rotation[0]), float(euler_rotation[1]), float(euler_rotation[2]))
self.xform.AddRotateXYZOp().Set(rotation)
# TODO: check to make sure scale and size are the same thing
def update_size(self, new_size):
size = tuple([float(x) for x in new_size])
self.xform.AddScaleOp().Set(value=size)
@@ -41,8 +62,7 @@ class USDGeom(object):
def update_color(self, new_color):
# new_color is the rgba (we extract first three)
rgba = [(float(x) for x in new_color[:3])]
color = self.prim.GetDisplayColorAttr()
color.Set(rgba)
color = self.prim.GetDisplayColorAttr().Set(rgba)
# TODO: create another method for transparency
def update_transparency(self, new_transparency):
@@ -70,10 +90,10 @@ class USDPlane(USDGeom):
self.prim = UsdGeom.Plane.Define(stage, plane_path)
self.ref = stage.GetPrimAtPath(plane_path)
def update_geom(self, new_geom):
self.update_pos(new_geom.pos)
self.update_size(new_geom.size)
self.update_color(new_geom.rgba)
def update_size(self, new_size):
self.prim.GetAxisAttr().Set("Z")
self.prim.GetWidthAttr().Set(float(new_size[0]))
self.prim.GetLengthAttr().Set(float(new_size[1]))
class USDSphere(USDGeom):
"""
@@ -83,8 +103,8 @@ class USDSphere(USDGeom):
sphere_count = 0
def __init__(self,
geom=None,
stage=None):
geom,
stage):
super().__init__(geom, stage)
self.type = 2
USDSphere.sphere_count += 1
@@ -94,11 +114,6 @@ class USDSphere(USDGeom):
self.prim = UsdGeom.Sphere.Define(stage, sphere_path)
self.ref = stage.GetPrimAtPath(sphere_path)
def update_geom(self, new_geom):
self.update_pos(new_geom.pos)
self.update_size(new_geom.size)
self.update_color(new_geom.rgba)
class USDCube(USDGeom):
"""
Stores information regarding a cube geom in USD
@@ -107,8 +122,8 @@ class USDCube(USDGeom):
cube_count = 0
def __init__(self,
geom=None,
stage=None):
geom,
stage):
super().__init__(geom, stage)
self.type = 6
USDCube.cube_count += 1
@@ -118,24 +133,76 @@ class USDCube(USDGeom):
self.prim = UsdGeom.Cube.Define(stage, cube_path)
self.ref = stage.GetPrimAtPath(cube_path)
class USDMesh(USDGeom):
"""
Stores information regarding a mesh geom in USD
"""
mesh_count = 0
def __init__(self,
mesh_idx,
geom,
stage,
model,
mesh_vertex_ranges,
mesh_face_ranges):
super().__init__(geom, stage)
self.type = 7
USDMesh.mesh_count += 1
xform_path = f'/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 = model.mesh_vert[mesh_vertex_ranges[mesh_idx]:mesh_vertex_ranges[mesh_idx+1]]
self.prim.GetPointsAttr().Set(self.vertices)
self.prim.GetFaceVertexCountsAttr().Set([3 for _ in range(model.mesh_facenum[mesh_idx])])
self.faces = model.mesh_face[mesh_face_ranges[mesh_idx]:mesh_face_ranges[mesh_idx+1]]
self.prim.GetFaceVertexIndicesAttr().Set(self.faces)
def update_geom(self, new_geom):
self.update_pos(new_geom.pos)
self.update_size(new_geom.size)
self.update_color(new_geom.rgba)
self.update_rotation(new_geom.mat)
def create_usd_geom(geom, stage):
geom_type = geom.type
if geom_type==USDGeomType.Plane.value:
return USDPlane(geom, stage)
elif geom_type==USDGeomType.Sphere.value:
return USDSphere(geom, stage)
elif geom_type==USDGeomType.Cube.value:
return USDCube(geom, stage)
# TODO: remove this, temporary
self.xform.AddScaleOp().Set(value=(10.0, 10.0, 10.0))
class USDLight(object):
def __init__(self):
pass
"""
Class for the created lights
"""
light_count = 0
def __init__(self,
stage):
self.stage = stage
USDLight.light_count += 1
xform_path = f'/Light_Xform_{USDLight.light_count}'
light_path = f'{xform_path}/Light_{USDLight.light_count}'
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdLux.SphereLight.Define(stage, light_path)
self.ref = stage.GetPrimAtPath(light_path)
def update_light(self, new_light):
pass
pos = tuple([float(x) for x in new_light.pos])
self.xform.AddTranslateOp().Set(pos)
# TODO attributes:
# - direction
# - intensity
# - exposure
# - radius
# - specular
self.prim.GetIntensityAttr().Set(5000);
python/mujoco/usd_renderer.py
+25 -5
View File
@@ -2,6 +2,7 @@
import mujoco
import mujoco.viewer as viewer
from mujoco.usd_component import *
from mujoco.usd_utilities import *
from pxr import Usd, UsdGeom
class USDRenderer(object):
@@ -47,13 +48,29 @@ class USDRenderer(object):
"""
Loads and initializes the necessary objects to render the scene
"""
if self.model.nmesh > 0:
mesh_vertex_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_vertnum)
mesh_face_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_facenum)
current_mesh_idx = 0
# 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):
self.usd_geoms.append(create_usd_geom(geoms[i], self.stage))
if geoms[i].type == USDGeomType.Mesh.value:
self.usd_geoms.append(USDMesh(current_mesh_idx,
geoms[i],
self.stage,
self.model,
mesh_vertex_ranges,
mesh_face_ranges))
current_mesh_idx += 1
else:
self.usd_geoms.append(create_usd_geom_primitive(geoms[i], self.stage))
# initializes an array to store all the lights in the scene
# populates with "empty" USDLight objects
@@ -61,7 +78,7 @@ class USDRenderer(object):
lights = self.scene.lights
self.nlight = self.scene.nlight
for i in range(self.nlight):
self.usd_lights.append(USDLight())
self.usd_lights.append(USDLight(self.stage))
def _update(self):
self._update_geoms()
@@ -83,9 +100,12 @@ class USDRenderer(object):
lights = self.scene.lights
nlight = self.scene.nlight
for i in range(nlight):
print(self.usd_lights[i])
self.usd_lights[i].update_light(lights[i])
def _update_camera(self):
"""
Updates the camera to match the current scene
"""
pass
def start_viewer(self):
@@ -94,11 +114,11 @@ class USDRenderer(object):
def render(self):
# should render the usd file given a particular renderer that
# works with USD files?
# works with USD files
# TODO: determine if this is valid functionality
pass
# TODO: remove later, this is only for debuggin purposes
# 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
+8
View File
@@ -0,0 +1,8 @@
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