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
texture_file = (
self.texture_files[
self.model.mat_texid[geom.matid][mujoco.mjtTextureRole.mjTEXROLE_RGB]
]
if geom.matid != -1
else None
geom_textures = (
[(self.texture_files[i], self.model.tex_type[i]) if i != -1 else None for i in self.model.mat_texid[geom.matid]]
if geom.matid != -1
else None
)
# handling meshes in our scene
@@ -268,7 +266,7 @@ class USDExporter:
obj_name=geom_name,
dataid=self.model.geom_dataid[geom.objid],
rgba=geom.rgba,
texture_file=texture_file,
geom_textures=geom_textures,
)
else:
# handling tendons in our scene
@@ -286,7 +284,7 @@ class USDExporter:
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
geom_textures=geom_textures,
)
# handling primitives in our scene
else:
@@ -302,7 +300,7 @@ class USDExporter:
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
geom_textures=geom_textures,
)
self.geom_names.add(geom_name)
python/mujoco/usd/objects.py
+58 -54
View File
@@ -16,7 +16,7 @@
import abc
import collections
from typing import Optional, Dict, Any
from typing import Optional, Dict, Any, Tuple
import mujoco
import mujoco.usd.shapes as shapes_module
@@ -55,14 +55,14 @@ class USDObject(abc.ABC):
geom: mujoco.MjvGeom,
obj_name: str,
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.model = model
self.geom = geom
self.obj_name = obj_name
self.rgba = rgba
self.texture_file = texture_file
self.geom_textures = geom_textures
self.xform_path = f"/World/Mesh_Xform_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
@@ -120,7 +120,7 @@ class USDObject(abc.ABC):
# setting the image texture attributes
image_shader.CreateIdAttr("UsdUVTexture")
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(
"sRGB"
@@ -225,9 +225,9 @@ class USDMesh(USDObject):
obj_name: str,
dataid: int,
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
@@ -243,15 +243,14 @@ class USDMesh(USDObject):
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# setting mesh uv properties
mesh_texcoord, mesh_facetexcoord = 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(mesh_facetexcoord.tolist()))
if self.texture_file:
if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
# setting mesh uv properties
mesh_texcoord, mesh_facetexcoord = 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(mesh_facetexcoord.tolist()))
self.attach_image_material(self.usd_mesh)
else:
self.attach_solid_material(self.usd_mesh)
@@ -320,9 +319,9 @@ class USDPrimitiveMesh(USDObject):
geom: mujoco.MjvGeom,
obj_name: str,
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.prim_mesh = self.generate_primitive_mesh()
@@ -337,31 +336,32 @@ class USDPrimitiveMesh(USDObject):
[3 for _ in range(mesh_facenum)]
)
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)
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)
else:
self.attach_solid_material(self.usd_mesh)
def generate_primitive_mesh(self):
"""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())
return prim_mesh
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]
@@ -400,9 +400,9 @@ class USDTendon(USDObject):
geom: mujoco.MjvGeom,
obj_name: str,
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.tendon_parts = self.generate_primitive_mesh()
@@ -432,33 +432,36 @@ class USDTendon(USDObject):
part_geometry["mesh_face"]
)
# setting uv properties for each of the parts in the tendon
part_uv_geometries = self._get_uv_geometry()
for name, part_uv_geometry in part_uv_geometries.items():
self.texcoords = UsdGeom.PrimvarsAPI(
self.usd_refs[name]["usd_mesh"]
).CreatePrimvar(
"UVMap",
Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying,
)
self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
self.texcoords.SetIndices(
Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
)
for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"])
if self.texture_file:
self.attach_image_material(ref["usd_mesh"])
else:
self.attach_solid_material(ref["usd_mesh"])
if self.geom_textures and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB]:
# setting uv properties for each of the parts in the tendon
part_uv_geometries = self._get_uv_geometry()
for name, part_uv_geometry in part_uv_geometries.items():
self.texcoords = UsdGeom.PrimvarsAPI(
self.usd_refs[name]["usd_mesh"]
).CreatePrimvar(
"UVMap",
Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying,
)
self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
self.texcoords.SetIndices(
Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
)
for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"])
self.attach_image_material(ref["usd_mesh"])
else:
for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"])
self.attach_solid_material(ref["usd_mesh"])
def generate_primitive_mesh(self):
"""Generates the tendon mesh using primitives."""
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:
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())
mesh_parts[mesh_name] = prim_mesh
return mesh_parts
@@ -466,6 +469,7 @@ class USDTendon(USDObject):
def _get_uv_geometry(self):
part_uv_geometries = collections.defaultdict(dict)
for name, mesh in self.tendon_parts.items():
assert mesh.triangle_uvs is not None
mesh_texcoord = np.array(mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(mesh.triangles)
part_uv_geometries[name] = {
python/mujoco/usd/shapes.py
+44 -18
View File
@@ -14,17 +14,24 @@
# ==============================================================================
"""Built-in shapes for USD exporter."""
from typing import Dict, Any, Tuple
from typing import Dict, Any, Tuple, Optional
import mujoco
import numpy as np
def get_triangle_uvs(vertices: np.array, triangles: np.array):
# (jabhi) asusming 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
def get_triangle_uvs(
vertices: np.array,
triangles: np.array,
texture_type: Optional[mujoco.mjtTexture]
):
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():
""" Store UV and geometry information for a primitve mesh
@@ -39,7 +46,11 @@ class TriangleMesh():
@classmethod
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],
[width, 0.0, 0.0],
@@ -63,13 +74,16 @@ class TriangleMesh():
[0, 4, 1],
[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)
@classmethod
def create_sphere(
cls, radius: float, resolution: int
cls,
radius: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
):
vertices = []
triangles = []
@@ -93,13 +107,16 @@ class TriangleMesh():
vertices = np.array(vertices)
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)
@classmethod
def create_hemisphere(
cls, radius: float, resolution: int
cls,
radius: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
):
vertices = []
triangles = []
@@ -128,13 +145,17 @@ class TriangleMesh():
vertices = np.array(vertices)
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)
@classmethod
def create_cylinder(
cls, radius: float, height: float, resolution: int
cls,
radius: float,
height: float,
texture_type: Optional[mujoco.mjtTexture],
resolution: int
):
vertices = []
triangles = []
@@ -166,7 +187,7 @@ class TriangleMesh():
vertices = np.array(vertices)
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)
@@ -189,7 +210,7 @@ class TriangleMesh():
new_vertices = np.vstack((self.vertices, other.vertices))
other_triangles = other.triangles + len(self.vertices)
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)
raise TypeError(f"Cannot add TriangleMesh with {type(other)}")
@@ -282,7 +303,8 @@ def mesh_config_generator(
def mesh_factory(
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."""
assert "name" in mesh_config
@@ -300,22 +322,26 @@ def mesh_factory(
prim_mesh = TriangleMesh.create_box(
width=mesh_config[shape]["width"],
height=mesh_config[shape]["height"],
depth=mesh_config[shape]["depth"]
depth=mesh_config[shape]["depth"],
texture_type=texture_type
)
elif "hemisphere" in shape:
prim_mesh = TriangleMesh.create_hemisphere(
radius=mesh_config[shape]["radius"],
texture_type=texture_type,
resolution=resolution
)
elif "sphere" in shape:
prim_mesh = TriangleMesh.create_sphere(
radius=mesh_config[shape]["radius"],
texture_type=texture_type,
resolution=resolution
)
elif "cylinder" in shape:
prim_mesh = TriangleMesh.create_cylinder(
radius=mesh_config[shape]["radius"],
height=mesh_config[shape]["height"],
texture_type=texture_type,
resolution=resolution
)
else: