diff --git a/python/dist/mujoco-2.3.8.tar.gz b/python/dist/mujoco-2.3.8.tar.gz index 1730bbc6..8b659e93 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 a5fd033e..7956b82d 100644 --- a/python/mujoco/usd_component.py +++ b/python/mujoco/usd_component.py @@ -6,12 +6,19 @@ from mujoco.usd_utilities import * from pxr import Usd, UsdGeom, UsdLux, UsdShade, Vt, Gf, Sdf from scipy.spatial.transform import Rotation as R +# TODO: clean this up and remove the if statements def create_usd_geom_primitive(geom, stage, texture_file): geom_type = geom.type if geom_type==USDGeomType.Plane.value: return USDPlane(geom, stage, texture_file) elif geom_type==USDGeomType.Sphere.value: return USDSphere(geom, stage, texture_file) + elif geom_type==USDGeomType.Capsule.value: + return USDCapsule(geom, stage, texture_file) + elif geom_type==USDGeomType.Cylinder.value: + return USDCylinder(geom, stage, texture_file) + elif geom_type==USDGeomType.Cube.value: + return USDCube(geom, stage, texture_file) elif geom_type==USDGeomType.Cube.value: return USDCube(geom, stage, texture_file) @@ -22,7 +29,11 @@ class USDGeomType(Enum): https://mujoco.readthedocs.io/en/latest/APIreference/APItypes.html#mjtgeom """ Plane = 0 + # Hfield = 1 Sphere = 2 + Capsule = 3 + # Ellipsoid = 4 + Cylinder = 5 Cube = 6 Mesh = 7 @@ -35,18 +46,40 @@ class USDGeom(object): stage=None, texture_file=None): self.geom = geom - self.stage = stage # TODO: remove, not being used + self.stage = stage self.texture_file = texture_file self.type = None self.xform = None self.prim = None self.ref = None + + def set_texture(self): + if self.texture_file: + mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(self.texture_file))[0]}") + mtl = UsdShade.Material.Define(self.stage, mtl_path) + shader = UsdShade.Shader.Define(self.stage, mtl_path.AppendPath("Shader")) + shader.CreateIdAttr("UsdPreviewSurface") + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0)) + shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5) + shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0) + + diffuse_tx = UsdShade.Shader.Define(self.stage, mtl_path.AppendPath("DiffuseColorTx")) + diffuse_tx.CreateIdAttr('UsdUVTexture') + + diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(self.texture_file) + diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3) + shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb') + mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface") + + self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) + UsdShade.MaterialBindingAPI(self.prim).Bind(mtl) def update_geom(self, new_geom): self.update_pos(new_geom.pos) self.update_rotation(new_geom.mat) self.update_size(new_geom.size) self.update_color(new_geom.rgba) + self.update_transparency(new_geom.rgba[3]) def update_pos(self, new_pos): pos = tuple([float(x) for x in new_pos]) @@ -66,11 +99,10 @@ 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().Set(rgba) + self.prim.GetDisplayColorAttr().Set(rgba) - # TODO: create another method for transparency def update_transparency(self, new_transparency): - pass + self.prim.GetDisplayOpacityAttr().Set(new_transparency) def __str__(self): return f'type = {self.type} \ngeom = {self.geom}' @@ -89,12 +121,14 @@ class USDPlane(USDGeom): super().__init__(geom, stage, texture_file) self.type = 0 USDPlane.plane_count += 1 - xform_path = f'/Plane_Xform_{USDPlane.plane_count}' + xform_path = f'/World/Plane_Xform_{USDPlane.plane_count}' plane_path = f'{xform_path}/Plane_{USDPlane.plane_count}' self.xform = UsdGeom.Xform.Define(stage, xform_path) self.prim = UsdGeom.Plane.Define(stage, plane_path) self.ref = stage.GetPrimAtPath(plane_path) + self.set_texture() + def update_size(self, new_size): self.prim.GetAxisAttr().Set("Z") self.prim.GetWidthAttr().Set(float(new_size[0])) @@ -114,12 +148,80 @@ class USDSphere(USDGeom): super().__init__(geom, stage, texture_file) self.type = 2 USDSphere.sphere_count += 1 - xform_path = f'/Sphere_Xform_{USDSphere.sphere_count}' + xform_path = f'/World/Sphere_Xform_{USDSphere.sphere_count}' sphere_path = f'{xform_path}/Sphere_{USDSphere.sphere_count}' self.xform = UsdGeom.Xform.Define(stage, xform_path) self.prim = UsdGeom.Sphere.Define(stage, sphere_path) self.ref = stage.GetPrimAtPath(sphere_path) + self.set_texture() + +class USDCapsule(USDGeom): + """ + Stores information regarding a capsule geom in USD + """ + + capsule_count = 0 + + def __init__(self, + geom=None, + stage=None, + texture_file=None): + super().__init__(geom, stage, texture_file) + self.type = 3 + USDCapsule.capsule_count += 1 + xform_path = f'/World/Capsule_Xform_{USDCapsule.capsule_count}' + capsule_path = f'{xform_path}/Capsule_{USDCapsule.capsule_count}' + self.xform = UsdGeom.Xform.Define(stage, xform_path) + self.prim = UsdGeom.Capsule.Define(stage, sphere_path) + self.ref = stage.GetPrimAtPath(capsule_path) + + self.set_texture() + +class USDCylinder(USDGeom): + """ + Stores information regarding a capsule geom in USD + """ + + cylinder_count = 0 + + def __init__(self, + geom=None, + stage=None, + texture_file=None): + super().__init__(geom, stage, texture_file) + self.type = 5 + USDCylinder.cylinder_count += 1 + xform_path = f'/World/Cylinder_Xform_{USDCylinder.cylinder_count}' + cylinder_path = f'{xform_path}/Cylinder_{USDCylinder.cylinder_count}' + self.xform = UsdGeom.Xform.Define(stage, xform_path) + self.prim = UsdGeom.Cylinder.Define(stage, cylinder_path) + self.ref = stage.GetPrimAtPath(cylinder_path) + + self.set_texture() + +class USDSphere(USDGeom): + """ + Stores information regarding a sphere geom in USD + """ + + sphere_count = 0 + + def __init__(self, + geom=None, + stage=None, + texture_file=None): + super().__init__(geom, stage, texture_file) + self.type = 2 + USDSphere.sphere_count += 1 + xform_path = f'/World/Sphere_Xform_{USDSphere.sphere_count}' + sphere_path = f'{xform_path}/Sphere_{USDSphere.sphere_count}' + self.xform = UsdGeom.Xform.Define(stage, xform_path) + self.prim = UsdGeom.Sphere.Define(stage, sphere_path) + self.ref = stage.GetPrimAtPath(sphere_path) + + self.set_texture() + class USDCube(USDGeom): """ Stores information regarding a cube geom in USD @@ -134,31 +236,13 @@ class USDCube(USDGeom): super().__init__(geom, stage, texture_file) self.type = 6 USDCube.cube_count += 1 - xform_path = f'/Cube_Xform_{USDCube.cube_count}' + xform_path = f'/World/Cube_Xform_{USDCube.cube_count}' cube_path = f'{xform_path}/Cube_{USDCube.cube_count}' self.xform = UsdGeom.Xform.Define(stage, xform_path) self.prim = UsdGeom.Cube.Define(stage, cube_path) self.ref = stage.GetPrimAtPath(cube_path) - if texture_file: - mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(texture_file))[0]}") - mtl = UsdShade.Material.Define(stage, mtl_path) - shader = UsdShade.Shader.Define(stage, mtl_path.AppendPath("Shader")) - shader.CreateIdAttr("UsdPreviewSurface") - shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0)) - shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5) - shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0) - - diffuse_tx = UsdShade.Shader.Define(stage,mtl_path.AppendPath("DiffuseColorTx")) - diffuse_tx.CreateIdAttr('UsdUVTexture') - - diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(texture_file) - diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3) - shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb') - mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface") - - self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI) - UsdShade.MaterialBindingAPI(self.prim).Bind(mtl) + self.set_texture() class USDMesh(USDGeom): """ @@ -186,7 +270,7 @@ class USDMesh(USDGeom): mesh_face = model.mesh_face[mesh_face_adr_from:mesh_face_adr_to] self.type = 7 - xform_path = f'/Mesh_Xform_{USDMesh.mesh_count}' + xform_path = f'/World/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) @@ -243,7 +327,6 @@ class USDMesh(USDGeom): USDMesh.mesh_count += 1 - def update_geom(self, new_geom): self.update_pos(new_geom.pos) self.update_rotation(new_geom.mat) @@ -264,7 +347,7 @@ class USDLight(object): stage): self.stage = stage USDLight.light_count += 1 - xform_path = f'/Light_Xform_{USDLight.light_count}' + xform_path = f'/World/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) @@ -281,7 +364,36 @@ class USDLight(object): # - radius # - specular - self.prim.GetIntensityAttr().Set(30000); + self.prim.GetIntensityAttr().Set(15000); + + +class USDCamera(object): + """ + Class for created cameras + """ + + camera_count = 0 + + def __init__(self, + stage): + self.stage = stage + USDCamera.camera_count += 1 + camera_path = f'/World/Camera_{USDCamera.camera_count}' + self.prim = UsdGeom.Camera.Define(stage, camera_path) + self.ref = stage.GetPrimAtPath(camera_path) + + def update_camera(self, new_camera): + # print("---- Updating camera in USD ----") + xformAPI = UsdGeom.XformCommonAPI(self.prim) + + # hardcoded values for Robosuite testing and prototype + # TODO: use actual camera values + xformAPI.SetTranslate((2.25, 0.0, 2.1)) + xformAPI.SetRotate((65.0, 0.0, 90.0)) + xformAPI.SetScale((1, 1, 1)) + + self.prim.CreateFocalLengthAttr().Set(24) + self.prim.CreateFocusDistanceAttr().Set(400) diff --git a/python/mujoco/usd_renderer.py b/python/mujoco/usd_renderer.py index 01b409ed..8cf4dd70 100644 --- a/python/mujoco/usd_renderer.py +++ b/python/mujoco/usd_renderer.py @@ -5,6 +5,8 @@ from mujoco.usd_component import * from mujoco.usd_utilities import * from pxr import Usd, UsdGeom +from mujoco import _structs + from PIL import Image as im from PIL import ImageOps @@ -15,7 +17,8 @@ class USDRenderer(object): def __init__(self, model, height=480, - width=480): + width=480, + geom_groups=[1,1,1,1,1,1]): self.model = model self.data = None self.renderer = mujoco.Renderer(model, height, width) @@ -25,6 +28,10 @@ class USDRenderer(object): self.stage = Usd.Stage.CreateNew('usd_stage.usda') UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z) + # TODO: maybe change where we initialize this? + self.scene_option = _structs.MjvOption() + self.scene_option.geomgroup = geom_groups + @property def usd(self): return self.stage.GetRootLayer().ExportToString() @@ -37,7 +44,7 @@ class USDRenderer(object): self.stage.GetRootLayer().Save() def update_scene(self, data): - self.renderer.update_scene(data) + self.renderer.update_scene(data, scene_option=self.scene_option) self.data = data if not self.loaded_scene_info: @@ -53,15 +60,9 @@ class USDRenderer(object): Loads and initializes the necessary objects to render the scene """ - # TODO: remove these and replace by reading directly from model - # 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) - # mesh_texcoord_ranges = get_mesh_ranges(self.model.nmesh, self.model.mesh_texcoordnum) - # mesh_facetexcoord_ranges = get_facetexcoord_ranges(self.model.nmesh, self.model.mesh_facenum) - # create and load the texture files # iterate through all the textures and build list of tex_rgb ranges + # TODO: remove once added to mujoco data_adr = 0 texture_files = [] for texid in range(self.model.ntex): @@ -85,20 +86,17 @@ class USDRenderer(object): self.ngeom = self.scene.ngeom for i in range(self.ngeom): geom = geoms[i] - if geom.category == 1: - self.usd_geoms.append(None) - continue if geom.texid == -1: texture_file = None else: texture_file = texture_files[geom.texid] if geom.type == USDGeomType.Mesh.value: - self.usd_geoms.append(USDMesh(self.model.geom_dataid[i], - geom, - self.stage, - self.model, - texture_file)) + self.usd_geoms.append(USDMesh(self.model.geom_dataid[geom.objid], + geom, + self.stage, + self.model, + texture_file)) else: self.usd_geoms.append(create_usd_geom_primitive(geom, self.stage, @@ -112,10 +110,17 @@ class USDRenderer(object): for i in range(self.nlight): self.usd_lights.append(USDLight(self.stage)) + # initializes an array to store all the cameras in the scene + # populates with "empty" USDCamera objects + self.usd_cameras = [] + cameras = self.scene.camera + for i in range(2): + self.usd_cameras.append(USDCamera(self.stage)) + def _update(self): self._update_geoms() self._update_lights() - self._update_camera() + self._update_cameras() def _update_geoms(self): """ @@ -135,11 +140,13 @@ class USDRenderer(object): for i in range(nlight): self.usd_lights[i].update_light(lights[i]) - def _update_camera(self): + def _update_cameras(self): """ Updates the camera to match the current scene """ - pass + cameras = self.scene.camera + for i in range(2): + self.usd_cameras[i].update_camera(cameras[i]) def start_viewer(self): if self.data: