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

Merge pull request #1509 from awesome-aj0123:camera-bug-fixes

Browse Source 
PiperOrigin-RevId: 617864822
Change-Id: If8bca106ecf725f9a7103eb3e009af901d6495a4
This commit is contained in:
Copybara-Service
2024-03-21 09:18:04 -07:00
parent c3fb8e7f35 7013c288bd
commit 67982c4961
2 changed files with 67 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/usd/component.py
+51 -49
View File
@@ -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)
python/mujoco/usd/exporter.py
+16 -15
View File
@@ -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)