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

Adding UV maps only when texture present

Browse Source
This commit is contained in:
Abhishek Joshi
2024-07-21 11:14:00 -05:00
parent 5abdeb8a71
commit d1ff008a4f
3 changed files with 109 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/usd/exporter.py
+7 -9
View File
@@ -251,12 +251,10 @@ class USDExporter:
assert geom_name not in self.geom_names assert geom_name not in self.geom_names
texture_file = ( geom_textures = (
self.texture_files[ [(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]]
self.model.mat_texid[geom.matid][mujoco.mjtTextureRole.mjTEXROLE_RGB] if geom.matid != -1
] else None
if geom.matid != -1
else None
) )
# handling meshes in our scene # handling meshes in our scene
@@ -268,7 +266,7 @@ class USDExporter:
obj_name=geom_name, obj_name=geom_name,
dataid=self.model.geom_dataid[geom.objid], dataid=self.model.geom_dataid[geom.objid],
rgba=geom.rgba, rgba=geom.rgba,
texture_file=texture_file, geom_textures=geom_textures,
) )
else: else:
# handling tendons in our scene # handling tendons in our scene
@@ -286,7 +284,7 @@ class USDExporter:
geom=geom, geom=geom,
obj_name=geom_name, obj_name=geom_name,
rgba=geom.rgba, rgba=geom.rgba,
texture_file=texture_file, geom_textures=geom_textures,
) )
# handling primitives in our scene # handling primitives in our scene
else: else:
@@ -302,7 +300,7 @@ class USDExporter:
geom=geom, geom=geom,
obj_name=geom_name, obj_name=geom_name,
rgba=geom.rgba, rgba=geom.rgba,
texture_file=texture_file, geom_textures=geom_textures,
) )
self.geom_names.add(geom_name) self.geom_names.add(geom_name)
python/mujoco/usd/objects.py
+58 -54
View File
@@ -16,7 +16,7 @@
import abc import abc
import collections import collections
from typing import Optional, Dict, Any from typing import Optional, Dict, Any, Tuple
import mujoco import mujoco
import mujoco.usd.shapes as shapes_module import mujoco.usd.shapes as shapes_module
@@ -55,14 +55,14 @@ class USDObject(abc.ABC):
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
obj_name: str, obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]), rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None, geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
): ):
self.stage = stage self.stage = stage
self.model = model self.model = model
self.geom = geom self.geom = geom
self.obj_name = obj_name self.obj_name = obj_name
self.rgba = rgba self.rgba = rgba
self.texture_file = texture_file self.geom_textures = geom_textures
self.xform_path = f"/World/Mesh_Xform_{obj_name}" self.xform_path = f"/World/Mesh_Xform_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path) self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
@@ -120,7 +120,7 @@ class USDObject(abc.ABC):
# setting the image texture attributes # setting the image texture attributes
image_shader.CreateIdAttr("UsdUVTexture") image_shader.CreateIdAttr("UsdUVTexture")
image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set( image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
self.texture_file self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB][0]
) )
image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set( image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set(
"sRGB" "sRGB"
@@ -225,9 +225,9 @@ class USDMesh(USDObject):
obj_name: str, obj_name: str,
dataid: int, dataid: int,
rgba: np.ndarray = np.array([1, 1, 1, 1]), rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None, geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
): ):
super().__init__(stage, model, geom, obj_name, rgba, texture_file) super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
self.dataid = dataid self.dataid = dataid
@@ -243,15 +243,14 @@ class USDMesh(USDObject):
) )
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# setting mesh uv properties if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry() # setting mesh uv properties
self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar( mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
"UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
) "UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
self.texcoords.Set(mesh_texcoord) )
self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist())) self.texcoords.Set(mesh_texcoord)
self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist()))
if self.texture_file:
self.attach_image_material(self.usd_mesh) self.attach_image_material(self.usd_mesh)
else: else:
self.attach_solid_material(self.usd_mesh) self.attach_solid_material(self.usd_mesh)
@@ -320,9 +319,9 @@ class USDPrimitiveMesh(USDObject):
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
obj_name: str, obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]), rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None, geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
): ):
super().__init__(stage, model, geom, obj_name, rgba, texture_file) super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
self.mesh_config = mesh_config self.mesh_config = mesh_config
self.prim_mesh = self.generate_primitive_mesh() self.prim_mesh = self.generate_primitive_mesh()
@@ -337,31 +336,32 @@ class USDPrimitiveMesh(USDObject):
[3 for _ in range(mesh_facenum)] [3 for _ in range(mesh_facenum)]
) )
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face) self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# setting mesh uv properties
mesh_texcoord, _ = self._get_uv_geometry()
self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
"UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
)
self.texcoords.Set(mesh_texcoord)
self.texcoords.SetIndices(Vt.IntArray(list(range(mesh_facenum * 3))))
self._set_refinement_properties(self.usd_prim) self._set_refinement_properties(self.usd_prim)
if self.texture_file: if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
# setting mesh uv properties
mesh_texcoord, _ = self._get_uv_geometry()
self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
"UVMap", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
)
self.texcoords.Set(mesh_texcoord)
self.texcoords.SetIndices(Vt.IntArray(list(range(mesh_facenum * 3))))
self.attach_image_material(self.usd_mesh) self.attach_image_material(self.usd_mesh)
else: else:
self.attach_solid_material(self.usd_mesh) self.attach_solid_material(self.usd_mesh)
def generate_primitive_mesh(self): def generate_primitive_mesh(self):
"""Generates the mesh for the primitive USD object.""" """Generates the mesh for the primitive USD object."""
_, prim_mesh = shapes_module.mesh_factory(self.mesh_config) geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] if self.geom_textures else None
texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
_, prim_mesh = shapes_module.mesh_factory(self.mesh_config, texture_type)
prim_mesh.translate(-prim_mesh.get_center()) prim_mesh.translate(-prim_mesh.get_center())
return prim_mesh return prim_mesh
def _get_uv_geometry(self): def _get_uv_geometry(self):
assert self.prim_mesh assert self.prim_mesh and self.prim_mesh.triangle_uvs is not None
s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid] s_scale, t_scale = self.model.mat_texrepeat[self.geom.matid]
@@ -400,9 +400,9 @@ class USDTendon(USDObject):
geom: mujoco.MjvGeom, geom: mujoco.MjvGeom,
obj_name: str, obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]), rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None, geom_textures: Optional[Tuple[str, mujoco.mjtTexture]] = None
): ):
super().__init__(stage, model, geom, obj_name, rgba, texture_file) super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
self.mesh_config = mesh_config self.mesh_config = mesh_config
self.tendon_parts = self.generate_primitive_mesh() self.tendon_parts = self.generate_primitive_mesh()
@@ -432,33 +432,36 @@ class USDTendon(USDObject):
part_geometry["mesh_face"] part_geometry["mesh_face"]
) )
# setting uv properties for each of the parts in the tendon if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
part_uv_geometries = self._get_uv_geometry() # setting uv properties for each of the parts in the tendon
for name, part_uv_geometry in part_uv_geometries.items(): part_uv_geometries = self._get_uv_geometry()
self.texcoords = UsdGeom.PrimvarsAPI( for name, part_uv_geometry in part_uv_geometries.items():
self.usd_refs[name]["usd_mesh"] self.texcoords = UsdGeom.PrimvarsAPI(
).CreatePrimvar( self.usd_refs[name]["usd_mesh"]
"UVMap", ).CreatePrimvar(
Sdf.ValueTypeNames.TexCoord2fArray, "UVMap",
UsdGeom.Tokens.faceVarying, Sdf.ValueTypeNames.TexCoord2fArray,
) UsdGeom.Tokens.faceVarying,
self.texcoords.Set(part_uv_geometry["mesh_texcoord"]) )
self.texcoords.SetIndices( self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3))) self.texcoords.SetIndices(
) Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
)
for _, ref in self.usd_refs.items(): for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"]) self._set_refinement_properties(ref["usd_prim"])
if self.texture_file: self.attach_image_material(ref["usd_mesh"])
self.attach_image_material(ref["usd_mesh"]) else:
else: for _, ref in self.usd_refs.items():
self.attach_solid_material(ref["usd_mesh"]) self._set_refinement_properties(ref["usd_prim"])
self.attach_solid_material(ref["usd_mesh"])
def generate_primitive_mesh(self): def generate_primitive_mesh(self):
"""Generates the tendon mesh using primitives.""" """Generates the tendon mesh using primitives."""
mesh_parts = {} mesh_parts = {}
geom_rgb_texture = self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] if self.geom_textures else None
texture_type = geom_rgb_texture[1] if geom_rgb_texture else None
for part_config in self.mesh_config: for part_config in self.mesh_config:
mesh_name, prim_mesh = shapes_module.mesh_factory(part_config) mesh_name, prim_mesh = shapes_module.mesh_factory(part_config, texture_type)
prim_mesh.translate(-prim_mesh.get_center()) prim_mesh.translate(-prim_mesh.get_center())
mesh_parts[mesh_name] = prim_mesh mesh_parts[mesh_name] = prim_mesh
return mesh_parts return mesh_parts
@@ -466,6 +469,7 @@ class USDTendon(USDObject):
def _get_uv_geometry(self): def _get_uv_geometry(self):
part_uv_geometries = collections.defaultdict(dict) part_uv_geometries = collections.defaultdict(dict)
for name, mesh in self.tendon_parts.items(): for name, mesh in self.tendon_parts.items():
assert mesh.triangle_uvs is not None
mesh_texcoord = np.array(mesh.triangle_uvs) mesh_texcoord = np.array(mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(mesh.triangles) mesh_facetexcoord = np.asarray(mesh.triangles)
part_uv_geometries[name] = { part_uv_geometries[name] = {
python/mujoco/usd/shapes.py
+44 -18
View File
@@ -14,17 +14,24 @@
# ============================================================================== # ==============================================================================
"""Built-in shapes for USD exporter.""" """Built-in shapes for USD exporter."""
from typing import Dict, Any, Tuple from typing import Dict, Any, Tuple, Optional
import mujoco import mujoco
import numpy as np import numpy as np
def get_triangle_uvs(vertices: np.array, triangles: np.array): def get_triangle_uvs(
# (jabhi) asusming all mappings are 2d mapping temporarily vertices: np.array,
triangle_uvs = np.array([ triangles: np.array,
[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)] texture_type: Optional[mujoco.mjtTexture]
) ):
return triangle_uvs if texture_type == None:
return None
# (jabhi) assuming all mappings are 2d mapping temporarily
triangle_uvs = np.array([
[vertices[i][0], vertices[i][1]] for i in np.nditer(triangles)]
)
return triangle_uvs
class TriangleMesh(): class TriangleMesh():
""" Store UV and geometry information for a primitve mesh """ Store UV and geometry information for a primitve mesh
@@ -39,7 +46,11 @@ class TriangleMesh():
@classmethod @classmethod
def create_box( def create_box(
cls, width: float, height: float, depth: float cls,
width: float,
height: float,
depth: float,
texture_type: Optional[mujoco.mjtTexture]
): ):
vertices = np.array([[0.0, 0.0, 0.0], vertices = np.array([[0.0, 0.0, 0.0],
[width, 0.0, 0.0], [width, 0.0, 0.0],
@@ -63,13 +74,16 @@ class TriangleMesh():
[0, 4, 1], [0, 4, 1],
[1, 4, 5]]) [1, 4, 5]])
triangle_uvs = get_triangle_uvs(vertices, triangles) triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs) return TriangleMesh(vertices, triangles, triangle_uvs)
@classmethod @classmethod
def create_sphere( def create_sphere(
cls, radius: float, resolution: int cls,
radius: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
): ):
vertices = [] vertices = []
triangles = [] triangles = []
@@ -93,13 +107,16 @@ class TriangleMesh():
vertices = np.array(vertices) vertices = np.array(vertices)
triangles = np.array(triangles) triangles = np.array(triangles)
triangle_uvs = get_triangle_uvs(vertices, triangles) triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs) return TriangleMesh(vertices, triangles, triangle_uvs)
@classmethod @classmethod
def create_hemisphere( def create_hemisphere(
cls, radius: float, resolution: int cls,
radius: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
): ):
vertices = [] vertices = []
triangles = [] triangles = []
@@ -128,13 +145,17 @@ class TriangleMesh():
vertices = np.array(vertices) vertices = np.array(vertices)
triangles = np.array(triangles) triangles = np.array(triangles)
triangle_uvs = get_triangle_uvs(vertices, triangles) triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs) return TriangleMesh(vertices, triangles, triangle_uvs)
@classmethod @classmethod
def create_cylinder( def create_cylinder(
cls, radius: float, height: float, resolution: int cls,
radius: float,
height: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
): ):
vertices = [] vertices = []
triangles = [] triangles = []
@@ -166,7 +187,7 @@ class TriangleMesh():
vertices = np.array(vertices) vertices = np.array(vertices)
triangles = np.array(triangles) triangles = np.array(triangles)
triangle_uvs = get_triangle_uvs(vertices, triangles) triangle_uvs = get_triangle_uvs(vertices, triangles, texture_type)
return TriangleMesh(vertices, triangles, triangle_uvs) return TriangleMesh(vertices, triangles, triangle_uvs)
@@ -189,7 +210,7 @@ class TriangleMesh():
new_vertices = np.vstack((self.vertices, other.vertices)) new_vertices = np.vstack((self.vertices, other.vertices))
other_triangles = other.triangles + len(self.vertices) other_triangles = other.triangles + len(self.vertices)
new_triangles = np.vstack((self.triangles, other_triangles)) new_triangles = np.vstack((self.triangles, other_triangles))
new_triangle_uvs = get_triangle_uvs(new_vertices, new_triangles) new_triangle_uvs = np.vstack((self.triangle_uvs, other.triangle_uvs))
return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs) return TriangleMesh(new_vertices, new_triangles, new_triangle_uvs)
raise TypeError(f"Cannot add TriangleMesh with {type(other)}") raise TypeError(f"Cannot add TriangleMesh with {type(other)}")
@@ -282,7 +303,8 @@ def mesh_config_generator(
def mesh_factory( def mesh_factory(
mesh_config: Dict[str, Any], mesh_config: Dict[str, Any],
resolution: int = 100, texture_type: Optional[mujoco.mjtTexture],
resolution: int = 10,
): ):
"""Generates a mesh given a config consisting of shapes.""" """Generates a mesh given a config consisting of shapes."""
assert "name" in mesh_config assert "name" in mesh_config
@@ -300,22 +322,26 @@ def mesh_factory(
prim_mesh = TriangleMesh.create_box( prim_mesh = TriangleMesh.create_box(
width=mesh_config[shape]["width"], width=mesh_config[shape]["width"],
height=mesh_config[shape]["height"], height=mesh_config[shape]["height"],
depth=mesh_config[shape]["depth"] depth=mesh_config[shape]["depth"],
texture_type=texture_type
) )
elif "hemisphere" in shape: elif "hemisphere" in shape:
prim_mesh = TriangleMesh.create_hemisphere( prim_mesh = TriangleMesh.create_hemisphere(
radius=mesh_config[shape]["radius"], radius=mesh_config[shape]["radius"],
texture_type=texture_type,
resolution=resolution resolution=resolution
) )
elif "sphere" in shape: elif "sphere" in shape:
prim_mesh = TriangleMesh.create_sphere( prim_mesh = TriangleMesh.create_sphere(
radius=mesh_config[shape]["radius"], radius=mesh_config[shape]["radius"],
texture_type=texture_type,
resolution=resolution resolution=resolution
) )
elif "cylinder" in shape: elif "cylinder" in shape:
prim_mesh = TriangleMesh.create_cylinder( prim_mesh = TriangleMesh.create_cylinder(
radius=mesh_config[shape]["radius"], radius=mesh_config[shape]["radius"],
height=mesh_config[shape]["height"], height=mesh_config[shape]["height"],
texture_type=texture_type,
resolution=resolution resolution=resolution
) )
else: else: