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

Including base colors for meshes

Browse Source
This commit is contained in:
Abhishek Joshi
2023-09-17 14:08:10 -05:00
parent 3f366ffd4c
commit deb8a8e2e2
3 changed files with 89 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/usd_component.py
+73 -41
View File
@@ -5,14 +5,14 @@ import mujoco
from pxr import Usd, UsdGeom, UsdLux, UsdShade, Vt, Gf, Sdf
from scipy.spatial.transform import Rotation as R
def create_usd_geom_primitive(geom, stage):
def create_usd_geom_primitive(geom, stage, texture_file):
geom_type = geom.type
if geom_type==USDGeomType.Plane.value:
return USDPlane(geom, stage)
return USDPlane(geom, stage, texture_file)
elif geom_type==USDGeomType.Sphere.value:
return USDSphere(geom, stage)
return USDSphere(geom, stage, texture_file)
elif geom_type==USDGeomType.Cube.value:
return USDCube(geom, stage)
return USDCube(geom, stage, texture_file)
class USDGeomType(Enum):
"""
@@ -31,14 +31,16 @@ class USDGeom(object):
"""
def __init__(self,
geom=None,
stage=None):
stage=None,
texture_file=None):
self.geom = geom
self.stage = stage # TODO: remove, not being used
self.texture_file = texture_file
self.type = None
self.xform = None
self.prim = None
self.ref = None
def update_geom(self, new_geom):
self.update_pos(new_geom.pos)
self.update_rotation(new_geom.mat)
@@ -81,8 +83,9 @@ class USDPlane(USDGeom):
def __init__(self,
geom=None,
stage=None):
super().__init__(geom, stage)
stage=None,
texture_file=None):
super().__init__(geom, stage, texture_file)
self.type = 0
USDPlane.plane_count += 1
xform_path = f'/Plane_Xform_{USDPlane.plane_count}'
@@ -104,9 +107,10 @@ class USDSphere(USDGeom):
sphere_count = 0
def __init__(self,
geom,
stage):
super().__init__(geom, stage)
geom=None,
stage=None,
texture_file=None):
super().__init__(geom, stage, texture_file)
self.type = 2
USDSphere.sphere_count += 1
xform_path = f'/Sphere_Xform_{USDSphere.sphere_count}'
@@ -123,9 +127,10 @@ class USDCube(USDGeom):
cube_count = 0
def __init__(self,
geom,
stage):
super().__init__(geom, stage)
geom=None,
stage=None,
texture_file=None):
super().__init__(geom, stage, texture_file)
self.type = 6
USDCube.cube_count += 1
xform_path = f'/Cube_Xform_{USDCube.cube_count}'
@@ -134,6 +139,26 @@ class USDCube(USDGeom):
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)
class USDMesh(USDGeom):
"""
Stores information regarding a mesh geom in USD
@@ -152,6 +177,9 @@ class USDMesh(USDGeom):
mesh_facetexcoord_ranges,
texture_file):
super().__init__(geom, stage)
assert mesh_idx != -1
self.type = 7
USDMesh.mesh_count += 1
xform_path = f'/Mesh_Xform_{USDMesh.mesh_count}'
@@ -159,7 +187,7 @@ class USDMesh(USDGeom):
self.xform = UsdGeom.Xform.Define(stage, xform_path)
self.prim = UsdGeom.Mesh.Define(stage, mesh_path)
self.ref = stage.GetPrimAtPath(mesh_path)
self.vertices = model.mesh_vert[mesh_vertex_ranges[mesh_idx]:mesh_vertex_ranges[mesh_idx+1]]
self.prim.GetPointsAttr().Set(self.vertices)
@@ -168,41 +196,45 @@ class USDMesh(USDGeom):
self.faces = model.mesh_face[mesh_face_ranges[mesh_idx]:mesh_face_ranges[mesh_idx+1]]
self.prim.GetFaceVertexIndicesAttr().Set(self.faces)
texid = geom.texid
texcoords = model.mesh_texcoord[mesh_texcoord_ranges[texid]:mesh_texcoord_ranges[texid+1]]
self.texture_file = texture_file
if texture_file:
facetexcoords = model.mesh_facetexcoord.flatten()
facetexcoords = facetexcoords[mesh_facetexcoord_ranges[mesh_idx]:mesh_facetexcoord_ranges[mesh_idx+1]]
self.texcoords = UsdGeom.PrimvarsAPI(self.prim).CreatePrimvar("st",
Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying)
texid = geom.texid
texcoords = model.mesh_texcoord[mesh_texcoord_ranges[texid]:mesh_texcoord_ranges[texid+1]]
self.texcoords.Set(texcoords)
self.texcoords.SetIndices(Vt.IntArray(facetexcoords.tolist()));
facetexcoords = model.mesh_facetexcoord.flatten()
facetexcoords = facetexcoords[mesh_facetexcoord_ranges[mesh_idx]:mesh_facetexcoord_ranges[mesh_idx+1]]
self.texcoords = UsdGeom.PrimvarsAPI(self.prim).CreatePrimvar("st",
Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying)
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)
self.texcoords.Set(texcoords)
self.texcoords.SetIndices(Vt.IntArray(facetexcoords.tolist()));
diffuse_tx = UsdShade.Shader.Define(stage,mtl_path.AppendPath("DiffuseColorTx"))
diffuse_tx.CreateIdAttr('UsdUVTexture')
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)
# TODO: don't hardcode the image file
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")
diffuse_tx = UsdShade.Shader.Define(stage,mtl_path.AppendPath("DiffuseColorTx"))
diffuse_tx.CreateIdAttr('UsdUVTexture')
self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.prim).Bind(mtl)
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)
def update_geom(self, new_geom):
self.update_pos(new_geom.pos)
self.update_rotation(new_geom.mat)
if not self.texture_file:
self.update_color(new_geom.rgba)
# TODO: remove this, temporary
# self.xform.AddScaleOp().Set(value=(10.0, 10.0, 10.0))
@@ -235,7 +267,7 @@ class USDLight(object):
# - radius
# - specular
self.prim.GetIntensityAttr().Set(5000);
self.prim.GetIntensityAttr().Set(30000);
python/mujoco/usd_renderer.py
+16 -8
View File
@@ -6,6 +6,7 @@ from mujoco.usd_utilities import *
from pxr import Usd, UsdGeom
from PIL import Image as im
from PIL import ImageOps
class USDRenderer(object):
"""
@@ -58,7 +59,6 @@ class USDRenderer(object):
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
data_adr = 0
@@ -70,20 +70,26 @@ class USDRenderer(object):
rgb = self.model.tex_rgb[data_adr:data_adr+pixels]
img = rgb.reshape(height, width, 3)
file_name = f'{texid}.png'
im.fromarray(img).save(file_name)
img = im.fromarray(img)
img = ImageOps.flip(img)
img.save(file_name)
texture_file = os.path.abspath(file_name)
texture_files.append(texture_file)
data_adr += pixels
current_mesh_idx = 0
# initializes an array to store all the geoms in the scene
# populates with "empty" USDGeom objects
self.usd_geoms = []
geoms = self.scene.geoms
self.ngeom = self.scene.ngeom
for i in range(self.ngeom):
if geoms[i].texid == -1:
texture_file = None
else:
texture_file = texture_files[geoms[i].texid]
if geoms[i].type == USDGeomType.Mesh.value:
self.usd_geoms.append(USDMesh(current_mesh_idx,
self.usd_geoms.append(USDMesh(self.model.geom_dataid[i],
geoms[i],
self.stage,
self.model,
@@ -91,10 +97,11 @@ class USDRenderer(object):
mesh_face_ranges,
mesh_texcoord_ranges,
mesh_facetexcoord_ranges,
texture_files[geoms[i].texid]))
current_mesh_idx += 1
texture_file))
else:
self.usd_geoms.append(create_usd_geom_primitive(geoms[i], self.stage))
self.usd_geoms.append(create_usd_geom_primitive(geoms[i],
self.stage,
texture_file))
# initializes an array to store all the lights in the scene
# populates with "empty" USDLight objects
@@ -115,7 +122,8 @@ class USDRenderer(object):
"""
geoms = self.scene.geoms
for i in range(self.ngeom):
self.usd_geoms[i].update_geom(geoms[i])
if self.usd_geoms[i] != None: # TODO: remove this once all primitives are added
self.usd_geoms[i].update_geom(geoms[i])
def _update_lights(self):
"""