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

Merge branch 'camera-bug-fixes' into demo_script

Browse Source
This commit is contained in:
Abhishek Joshi
2024-03-18 15:50:57 -05:00
parent 983837e830 7013c288bd
commit 765ba22cd9
2 changed files with 47 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/usd/component.py
+44 -44
View File
@@ -36,7 +36,7 @@ class USDMesh:
stage: Usd.Stage, stage: Usd.Stage,
model: mujoco.MjModel, model: mujoco.MjModel,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
dataid: int, dataid: int,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, texture_file: Optional[str] = None,
@@ -44,13 +44,13 @@ class USDMesh:
self.stage = stage self.stage = stage
self.model = model self.model = model
self.geom = geom self.geom = geom
self.objid = objid self.obj_name = obj_name
self.rgba = rgba self.rgba = rgba
self.dataid = dataid self.dataid = dataid
self.texture_file = texture_file self.texture_file = texture_file
xform_path = f"/World/Mesh_Xform_{objid}" xform_path = f"/World/Mesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/Mesh_{objid}" mesh_path = f"{xform_path}/Mesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path)
@@ -132,7 +132,7 @@ class USDMesh:
return mesh_vert, mesh_face, mesh_facenum return mesh_vert, mesh_face, mesh_facenum
def _attach_material(self): 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) mtl = UsdShade.Material.Define(self.stage, mtl_path)
if self.texture_file: if self.texture_file:
@@ -235,13 +235,13 @@ class USDPrimitiveMesh:
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, texture_file: Optional[str] = None,
): ):
self.stage = stage self.stage = stage
self.geom = geom self.geom = geom
self.objid = objid self.obj_name = obj_name
self.rgba = rgba self.rgba = rgba
self.texture_file = texture_file self.texture_file = texture_file
@@ -282,7 +282,7 @@ class USDPrimitiveMesh:
return mesh_vert, mesh_face, len(mesh_face) return mesh_vert, mesh_face, len(mesh_face)
def _attach_material(self): 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) mtl = UsdShade.Material.Define(self.stage, mtl_path)
if self.texture_file: if self.texture_file:
bsdf_shader = UsdShade.Shader.Define( bsdf_shader = UsdShade.Shader.Define(
@@ -383,13 +383,13 @@ class USDPrimitive:
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, texture_file: Optional[str] = None,
): ):
self.stage = stage self.stage = stage
self.geom = geom self.geom = geom
self.objid = objid self.obj_name = obj_name
self.rgba = rgba self.rgba = rgba
self.texture_file = texture_file self.texture_file = texture_file
@@ -401,7 +401,7 @@ class USDPrimitive:
self.usd_prim.GetAttribute("subdivisionScheme").Set("none") self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
def _attach_material(self): 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) mtl = UsdShade.Material.Define(self.stage, mtl_path)
if self.texture_file: if self.texture_file:
bsdf_shader = UsdShade.Shader.Define( bsdf_shader = UsdShade.Shader.Define(
@@ -500,15 +500,15 @@ class USDCapsule(USDPrimitive):
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, 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}" xform_path = f"/World/Capsule_Xform_{obj_name}"
capsule_path = f"{xform_path}/Capsule_{objid}" capsule_path = f"{xform_path}/Capsule_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_primitive_shape = UsdGeom.Capsule.Define(stage, capsule_path) self.usd_primitive_shape = UsdGeom.Capsule.Define(stage, capsule_path)
self.usd_prim = stage.GetPrimAtPath(capsule_path) self.usd_prim = stage.GetPrimAtPath(capsule_path)
@@ -536,15 +536,15 @@ class USDEllipsoid(USDPrimitive):
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, 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}" xform_path = f"/World/Ellipsoid_Xform_{obj_name}"
ellipsoid_path = f"{xform_path}/Ellipsoid_{objid}" ellipsoid_path = f"{xform_path}/Ellipsoid_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_primitive_shape = UsdGeom.Sphere.Define(stage, ellipsoid_path) self.usd_primitive_shape = UsdGeom.Sphere.Define(stage, ellipsoid_path)
self.usd_prim = stage.GetPrimAtPath(ellipsoid_path) self.usd_prim = stage.GetPrimAtPath(ellipsoid_path)
@@ -569,15 +569,15 @@ class USDCubeMesh(USDPrimitiveMesh):
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, 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}" xform_path = f"/World/CubeMesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/CubeMesh_{objid}" mesh_path = f"{xform_path}/CubeMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path)
@@ -623,15 +623,15 @@ class USDSphereMesh(USDPrimitiveMesh):
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, 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}" xform_path = f"/World/SphereMesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/SphereMesh_{objid}" mesh_path = f"{xform_path}/SphereMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path)
@@ -673,15 +673,15 @@ class USDCylinderMesh(USDPrimitiveMesh):
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, 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}" xform_path = f"/World/CylinderMesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/CylinderMesh_{objid}" mesh_path = f"{xform_path}/CylinderMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path)
@@ -725,15 +725,15 @@ class USDPlaneMesh(USDPrimitiveMesh):
self, self,
stage: Usd.Stage, stage: Usd.Stage,
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
objid: int, obj_name: str,
rgba: Tuple[int, ...] = (1, 1, 1, 1), rgba: Tuple[int, ...] = (1, 1, 1, 1),
texture_file: Optional[str] = None, 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}" xform_path = f"/World/Plane_Xform_{obj_name}"
plane_path = f"{xform_path}/PlaneMesh_{objid}" plane_path = f"{xform_path}/PlaneMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, plane_path) self.usd_mesh = UsdGeom.Mesh.Define(stage, plane_path)
self.usd_prim = stage.GetPrimAtPath(plane_path) self.usd_prim = stage.GetPrimAtPath(plane_path)
@@ -776,12 +776,12 @@ class USDPlaneMesh(USDPrimitiveMesh):
class USDSphereLight: class USDSphereLight:
def __init__( def __init__(
self, stage: Usd.Stage, objid: int, radius: Optional[float] = 0.3 self, stage: Usd.Stage, obj_name: str, radius: Optional[float] = 0.3
): ):
self.stage = stage self.stage = stage
xform_path = f"/World/Light_Xform_{objid}" xform_path = f"/World/Light_Xform_{obj_name}"
light_path = f"{xform_path}/Light_{objid}" light_path = f"{xform_path}/Light_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_light = UsdLux.SphereLight.Define(stage, light_path) self.usd_light = UsdLux.SphereLight.Define(stage, light_path)
self.usd_prim = stage.GetPrimAtPath(light_path) self.usd_prim = stage.GetPrimAtPath(light_path)
@@ -808,11 +808,11 @@ class USDSphereLight:
class USDDomeLight: class USDDomeLight:
def __init__(self, stage: Usd.Stage, objid: int): def __init__(self, stage: Usd.Stage, obj_name: str):
self.stage = stage self.stage = stage
xform_path = f"/World/Light_Xform_{objid}" xform_path = f"/World/Light_Xform_{obj_name}"
light_path = f"{xform_path}/Light_{objid}" light_path = f"{xform_path}/Light_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_light = UsdLux.DomeLight.Define(stage, light_path) self.usd_light = UsdLux.DomeLight.Define(stage, light_path)
self.usd_prim = stage.GetPrimAtPath(light_path) self.usd_prim = stage.GetPrimAtPath(light_path)
@@ -829,11 +829,11 @@ class USDDomeLight:
class USDCamera: class USDCamera:
def __init__(self, stage: Usd.Stage, objid: int): def __init__(self, stage: Usd.Stage, obj_name: str):
self.stage = stage self.stage = stage
xform_path = f"/World/Camera_Xform_{objid}" xform_path = f"/World/Camera_Xform_{obj_name}"
camera_path = f"{xform_path}/Camera_{objid}" camera_path = f"{xform_path}/Camera_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_camera = UsdGeom.Camera.Define(stage, camera_path) self.usd_camera = UsdGeom.Camera.Define(stage, camera_path)
self.usd_prim = stage.GetPrimAtPath(camera_path) self.usd_prim = stage.GetPrimAtPath(camera_path)
python/mujoco/usd/exporter.py
+3 -4
View File
@@ -352,12 +352,11 @@ class USDExporter:
frame=self.updates, frame=self.updates,
) )
print("done updating")
def _load_cameras(self): def _load_cameras(self):
self.usd_cameras = [] self.usd_cameras = []
for name in self.camera_names: if self.camera_names is not None:
self.usd_cameras.append(mujoco.usd.component.USDCamera(stage=self.stage, objid=name)) for name in self.camera_names:
self.usd_cameras.append(mujoco.usd.USDCamera(stage=self.stage, objid=name))
def _update_cameras( def _update_cameras(
self, self,