diff --git a/python/dist/mujoco-2.3.8.tar.gz b/python/dist/mujoco-2.3.8.tar.gz index 9f101d96..b73593af 100644 Binary files a/python/dist/mujoco-2.3.8.tar.gz and b/python/dist/mujoco-2.3.8.tar.gz differ diff --git a/python/mujoco/usd_component.py b/python/mujoco/usd_component.py index 7dee1993..732fb902 100644 --- a/python/mujoco/usd_component.py +++ b/python/mujoco/usd_component.py @@ -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); + + + + + + \ No newline at end of file diff --git a/python/mujoco/usd_renderer.py b/python/mujoco/usd_renderer.py index d17b9ef5..08b94fb5 100644 --- a/python/mujoco/usd_renderer.py +++ b/python/mujoco/usd_renderer.py @@ -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]) \ No newline at end of file diff --git a/python/mujoco/usd_utilities.py b/python/mujoco/usd_utilities.py new file mode 100644 index 00000000..6989c0b5 --- /dev/null +++ b/python/mujoco/usd_utilities.py @@ -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 \ No newline at end of file