diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py index c8a80f6f..044689ab 100644 --- a/python/mujoco/usd/component.py +++ b/python/mujoco/usd/component.py @@ -37,7 +37,7 @@ class USDMesh: stage: Usd.Stage, model: mujoco.MjModel, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, dataid: int, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, @@ -45,13 +45,13 @@ class USDMesh: self.stage = stage self.model = model self.geom = geom - self.objid = objid + self.obj_name = obj_name self.rgba = rgba self.dataid = dataid self.texture_file = texture_file - xform_path = f"/World/Mesh_Xform_{objid}" - mesh_path = f"{xform_path}/Mesh_{objid}" + xform_path = f"/World/Mesh_Xform_{obj_name}" + mesh_path = f"{xform_path}/Mesh_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path) @@ -133,7 +133,7 @@ class USDMesh: return mesh_vert, mesh_face, mesh_facenum def _attach_material(self): - mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}") + mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}") mtl = UsdShade.Material.Define(self.stage, mtl_path) if self.texture_file: @@ -236,13 +236,13 @@ class USDPrimitiveMesh: self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): self.stage = stage self.geom = geom - self.objid = objid + self.obj_name = obj_name self.rgba = rgba self.texture_file = texture_file @@ -283,7 +283,7 @@ class USDPrimitiveMesh: return mesh_vert, mesh_face, len(mesh_face) def _attach_material(self): - mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}") + mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}") mtl = UsdShade.Material.Define(self.stage, mtl_path) if self.texture_file: bsdf_shader = UsdShade.Shader.Define( @@ -384,25 +384,25 @@ class USDPrimitive: self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): self.stage = stage self.geom = geom - self.objid = objid + self.obj_name = obj_name self.rgba = rgba self.texture_file = texture_file self.usd_prim = Usd.Prim() self.usd_primitive_shape = Usd.PrimitiveShape() - self.transform_op = Gf.Matrix4d(1.) + self.transform_op = Usd.TransformOp() 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}") + mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}") mtl = UsdShade.Material.Define(self.stage, mtl_path) if self.texture_file: bsdf_shader = UsdShade.Shader.Define( @@ -482,7 +482,7 @@ class USDPrimitive: UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl) def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int): - transformation_mat = mujoco.usd.utils.create_transform_matrix( + transformation_mat = mujoco.usd_util.create_transform_matrix( rotation_matrix=mat, translation_vector=pos ).T self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame) @@ -501,15 +501,15 @@ class USDCapsule(USDPrimitive): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, objid, rgba, texture_file) + super().__init__(stage, geom, obj_name, rgba, texture_file) - xform_path = f"/World/Capsule_Xform_{objid}" - capsule_path = f"{xform_path}/Capsule_{objid}" + xform_path = f"/World/Capsule_Xform_{obj_name}" + capsule_path = f"{xform_path}/Capsule_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_primitive_shape = UsdGeom.Capsule.Define(stage, capsule_path) self.usd_prim = stage.GetPrimAtPath(capsule_path) @@ -537,15 +537,15 @@ class USDEllipsoid(USDPrimitive): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, objid, rgba, texture_file) + super().__init__(stage, geom, obj_name, rgba, texture_file) - xform_path = f"/World/Ellipsoid_Xform_{objid}" - ellipsoid_path = f"{xform_path}/Ellipsoid_{objid}" + xform_path = f"/World/Ellipsoid_Xform_{obj_name}" + ellipsoid_path = f"{xform_path}/Ellipsoid_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_primitive_shape = UsdGeom.Sphere.Define(stage, ellipsoid_path) self.usd_prim = stage.GetPrimAtPath(ellipsoid_path) @@ -570,15 +570,15 @@ class USDCubeMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, objid, rgba, texture_file) + super().__init__(stage, geom, obj_name, rgba, texture_file) - xform_path = f"/World/CubeMesh_Xform_{objid}" - mesh_path = f"{xform_path}/CubeMesh_{objid}" + xform_path = f"/World/CubeMesh_Xform_{obj_name}" + mesh_path = f"{xform_path}/CubeMesh_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path) @@ -622,21 +622,21 @@ class USDSphereMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, objid, rgba, texture_file) + super().__init__(stage, geom, obj_name, rgba, texture_file) - xform_path = f"/World/SphereMesh_Xform_{objid}" - mesh_path = f"{xform_path}/SphereMesh_{objid}" + xform_path = f"/World/SphereMesh_Xform_{obj_name}" + mesh_path = f"{xform_path}/SphereMesh_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path) self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere( - radius=float(self.geom.size[0]) + radius=float(self.geom.size[0]), create_uv_map=True ) self.prim_mesh.translate(-self.prim_mesh.get_center()) @@ -672,15 +672,15 @@ class USDCylinderMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, objid, rgba, texture_file) + super().__init__(stage, geom, obj_name, rgba, texture_file) - xform_path = f"/World/CylinderMesh_Xform_{objid}" - mesh_path = f"{xform_path}/CylinderMesh_{objid}" + xform_path = f"/World/CylinderMesh_Xform_{obj_name}" + mesh_path = f"{xform_path}/CylinderMesh_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path) @@ -688,6 +688,7 @@ class USDCylinderMesh(USDPrimitiveMesh): 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()) @@ -723,15 +724,15 @@ class USDPlaneMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: str, + obj_name: str, rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): - super().__init__(stage, geom, objid, rgba, texture_file) + super().__init__(stage, geom, obj_name, rgba, texture_file) - xform_path = f"/World/Plane_Xform_{objid}" - plane_path = f"{xform_path}/PlaneMesh_{objid}" + xform_path = f"/World/Plane_Xform_{obj_name}" + plane_path = f"{xform_path}/PlaneMesh_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, plane_path) self.usd_prim = stage.GetPrimAtPath(plane_path) @@ -740,6 +741,8 @@ class USDPlaneMesh(USDPrimitiveMesh): 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()) @@ -772,12 +775,12 @@ class USDPlaneMesh(USDPrimitiveMesh): class USDSphereLight: def __init__( - self, stage: Usd.Stage, objid: str, radius: Optional[float] = 0.3 + self, stage: Usd.Stage, obj_name: str, radius: Optional[float] = 0.3 ): self.stage = stage - xform_path = f"/World/Light_Xform_{objid}" - light_path = f"{xform_path}/Light_{objid}" + xform_path = f"/World/Light_Xform_{obj_name}" + light_path = f"{xform_path}/Light_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_light = UsdLux.SphereLight.Define(stage, light_path) self.usd_prim = stage.GetPrimAtPath(light_path) @@ -804,11 +807,11 @@ class USDSphereLight: class USDDomeLight: - def __init__(self, stage: Usd.Stage, objid: str): + def __init__(self, stage: Usd.Stage, obj_name: str): self.stage = stage - xform_path = f"/World/Light_Xform_{objid}" - light_path = f"{xform_path}/Light_{objid}" + xform_path = f"/World/Light_Xform_{obj_name}" + light_path = f"{xform_path}/Light_{obj_name}" 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) @@ -817,7 +820,6 @@ class USDDomeLight: 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())) @@ -825,11 +827,11 @@ class USDDomeLight: class USDCamera: - def __init__(self, stage: Usd.Stage, objid: str): + def __init__(self, stage: Usd.Stage, obj_name: str): self.stage = stage - xform_path = f"/World/Camera_Xform_{objid}" - camera_path = f"{xform_path}/Camera_{objid}" + xform_path = f"/World/Camera_Xform_{obj_name}" + camera_path = f"{xform_path}/Camera_{obj_name}" self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_camera = UsdGeom.Camera.Define(stage, camera_path) self.usd_prim = stage.GetPrimAtPath(camera_path) @@ -847,7 +849,7 @@ class USDCamera: def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int): - transformation_mat = mujoco.usd.utils.create_transform_matrix( + transformation_mat = mujoco.usd_util.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/exporter.py b/python/mujoco/usd/exporter.py index 4ec94bf6..33513ec2 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -240,7 +240,7 @@ class USDExporter: stage=self.stage, model=self.model, geom=geom, - objid=geom_name, + obj_name=geom_name, dataid=self.model.geom_dataid[geom.objid], rgba=geom.rgba, texture_file=texture_file, @@ -249,7 +249,7 @@ class USDExporter: usd_geom = component_module.USDPlaneMesh( stage=self.stage, geom=geom, - objid=geom_name, + obj_name=geom_name, rgba=geom.rgba, texture_file=texture_file, ) @@ -257,7 +257,7 @@ class USDExporter: usd_geom = component_module.USDSphereMesh( stage=self.stage, geom=geom, - objid=geom_name, + obj_name=geom_name, rgba=geom.rgba, texture_file=texture_file, ) @@ -265,7 +265,7 @@ class USDExporter: usd_geom = component_module.USDCapsule( stage=self.stage, geom=geom, - objid=geom_name, + obj_name=geom_name, rgba=geom.rgba, texture_file=texture_file, ) @@ -273,7 +273,7 @@ class USDExporter: usd_geom = component_module.USDEllipsoid( stage=self.stage, geom=geom, - objid=geom_name, + obj_name=geom_name, rgba=geom.rgba, texture_file=texture_file, ) @@ -281,7 +281,7 @@ class USDExporter: usd_geom = component_module.USDCylinderMesh( stage=self.stage, geom=geom, - objid=geom_name, + obj_name=geom_name, rgba=geom.rgba, texture_file=texture_file, ) @@ -289,7 +289,7 @@ class USDExporter: usd_geom = component_module.USDCubeMesh( stage=self.stage, geom=geom, - objid=geom_name, + obj_name=geom_name, rgba=geom.rgba, texture_file=texture_file, ) @@ -333,7 +333,7 @@ class USDExporter: light = self.scene.lights[i] if not np.allclose(light.pos, [0, 0, 0]): self.usd_lights.append - (component_module.USDSphereLight(stage=self.stage, objid=str(i))) + (component_module.USDSphereLight(stage=self.stage, obj_name=str(i))) else: self.usd_lights.append(None) @@ -354,14 +354,12 @@ class USDExporter: frame=self.updates, ) - print("done updating") - def _load_cameras(self): self.usd_cameras = [] if self.camera_names is not None: for name in self.camera_names: self.usd_cameras.append( - component_module.USDCamera(stage=self.stage, objid=name)) + component_module.USDCamera(stage=self.stage, obj_name=name)) def _update_cameras( self, @@ -377,7 +375,8 @@ class USDExporter: data, scene_option=scene_option, camera=camera_name ) - avg_camera = mujoco.mjv_averageCamera(self.scene.camera[0], self.scene.camera[1]) + avg_camera = mujoco.mjv_averageCamera( + self.scene.camera[0], self.scene.camera[1]) forward = avg_camera.forward up = avg_camera.up @@ -401,11 +400,12 @@ class USDExporter: ): if light_type == "sphere": - new_light = component_module.USDSphereLight(stage=self.stage, objid=str(objid), radius=radius) + new_light = component_module.USDSphereLight(stage=self.stage, obj_name=str(objid), radius=radius) new_light.update(pos=np.array(pos), intensity=intensity, color=color, frame=0) elif light_type == "dome": - new_light = component_module.USDDomeLight(stage=self.stage, objid=str(objid)) + new_light = component_module.USDDomeLight( + stage=self.stage, obj_name=str(objid)) new_light.update(intensity=intensity, color=color, frame=0) @@ -415,7 +415,8 @@ class USDExporter: rotation_xyz: List[float], objid: Optional[int] = 1, ): - new_camera = component_module.USDCamera(stage=self.stage, objid=str(objid)) + new_camera = component_module.USDCamera( + stage=self.stage, obj_name=str(objid)) r = scipy.spatial.transform.Rotation.from_euler( "xyz", rotation_xyz, degrees=True)