Visualizing specific geom groups
This commit is contained in:
Abhishek Joshi
+142
-30
@@ -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)
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user