chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 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/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] = {