Adds minor fix to handle normalmaps and blending

src/mjviser/conversions.py

@@ -22,22 +22,46 @@ def _get_texture_id(mj_model: mujoco.MjModel, matid: int) -> int:
   return texid
 
 
-def _extract_texture_image(mj_model: mujoco.MjModel, texid: int) -> Image.Image | None:
-  """Extract a 2D texture as a PIL Image, or None for unsupported types."""
+def _get_texture_normalmap_id(mj_model: mujoco.MjModel, matid: int) -> int:
+  """Returns the normalmap texture ID for a material, or -1."""
+  return int(mj_model.mat_texid[matid, int(mujoco.mjtTextureRole.mjTEXROLE_NORMAL)])
+
+
+def _has_alpha(image: Image.Image) -> bool:
+  """Return True if the image has an alpha channel with any transparent pixel."""
+  if image.mode != "RGBA":
+    return False
+  return bool(np.asarray(image.getchannel("A")).min() < 255)
+
+
+def _extract_texture_image(
+  mj_model: mujoco.MjModel, texid: int, flip: bool = True
+) -> Image.Image | None:
+  """Extract a 2D texture as a PIL Image, or None for unsupported types.
+
+  Color textures use MuJoCo's bottom-left origin while GLTF expects top-left,
+  so they are flipped vertically by default. Normal maps are stored without
+  that flip (they are typically authored with `vflip="false"` while color maps
+  use `vflip="true"`), so pass flip=False to keep them aligned with the albedo
+  texture and the shared UVs.
+  """
   w = mj_model.tex_width[texid]
   h = mj_model.tex_height[texid]
   nc = mj_model.tex_nchannel[texid]
   adr = mj_model.tex_adr[texid]
   data = mj_model.tex_data[adr : adr + w * h * nc]
 
-  # MuJoCo uses bottom-left origin; GLTF expects top-left.
   if nc == 1:
-    arr = np.flipud(data.reshape(h, w))
-    return Image.fromarray(arr.astype(np.uint8), mode="L")
+    arr = data.reshape(h, w)
   elif nc in (3, 4):
-    arr = np.flipud(data.reshape(h, w, nc))
-    return Image.fromarray(arr.astype(np.uint8))
-  return None
+    arr = data.reshape(h, w, nc)
+  else:
+    return None
+
+  if flip:
+    arr = np.flipud(arr)
+  mode = "L" if nc == 1 else None
+  return Image.fromarray(arr.astype(np.uint8), mode=mode)
 
 
 def _resolve_flat_rgba(mj_model: mujoco.MjModel, geom_idx: int) -> np.ndarray:
@@ -190,19 +214,32 @@ def mujoco_mesh_to_trimesh(mj_model: mujoco.MjModel, geom_idx: int) -> trimesh.T
 
   # Path 1: mesh has UVs and material has a 2D texture.
   if uvs is not None and matid >= 0:
-    texid = _get_texture_id(mj_model, matid)
-    if texid >= 0:
-      image = _extract_texture_image(mj_model, texid)
-      if image is not None:
-        rgba = mj_model.mat_rgba[matid]
-        material = trimesh.visual.material.PBRMaterial(
-          baseColorFactor=rgba,
-          baseColorTexture=image,
-          metallicFactor=0.0,
-          roughnessFactor=1.0,
-        )
-        mesh.visual = trimesh.visual.TextureVisuals(uv=uvs, material=material)
-        return mesh
+    texid_albedo = _get_texture_id(mj_model, matid)
+    texid_normalmap = _get_texture_normalmap_id(mj_model, matid)
+
+    image_albedo: Image.Image | None = None
+    image_normalmap: Image.Image | None = None
+    if texid_albedo >= 0:
+      image_albedo = _extract_texture_image(mj_model, texid_albedo)
+    if texid_normalmap >= 0:
+      image_normalmap = _extract_texture_image(mj_model, texid_normalmap, flip=False)
+
+    if image_albedo is not None:
+      rgba = mj_model.mat_rgba[matid]
+      geom_rgba = mj_model.geom_rgba[geom_idx]
+      # Blend when the geom/material is translucent or the texture itself has
+      # transparent pixels (e.g. cut-out decals stored in the alpha channel).
+      use_blending = rgba[-1] < 0.99 or geom_rgba[-1] < 0.99 or _has_alpha(image_albedo)
+      material = trimesh.visual.material.PBRMaterial(
+        baseColorFactor=rgba,
+        baseColorTexture=image_albedo,
+        metallicFactor=0.0,
+        roughnessFactor=1.0,
+        normalTexture=image_normalmap,
+        alphaMode="BLEND" if use_blending else "OPAQUE",
+      )
+      mesh.visual = trimesh.visual.TextureVisuals(uv=uvs, material=material)
+      return mesh
 
   # Path 2: no UVs, try cube map projection.
   if uvs is None and matid >= 0:

tests/conftest.py

@@ -83,11 +83,41 @@ def simple_data(simple_model):
 """
 
 
+_TEXTURED_XML = """
+<mujoco>
+  <asset>
+    <texture name="albedo" type="2d" builtin="checker" width="64" height="64"
+             rgb1="0.8 0.2 0.2" rgb2="0.2 0.2 0.8"/>
+    <texture name="nrm" type="2d" builtin="flat" width="64" height="64"
+             rgb1="0.5 0.5 1" rgb2="0.5 0.5 1"/>
+    <material name="mat_normal">
+      <layer texture="albedo" role="rgb"/>
+      <layer texture="nrm" role="normal"/>
+    </material>
+    <material name="mat_plain" texture="albedo"/>
+    <mesh name="quad" vertex="0 0 0  1 0 0  1 1 0  0 1 0.02"
+          texcoord="0 0  1 0  1 1  0 1" face="0 1 2  0 2 3"/>
+  </asset>
+  <worldbody>
+    <geom name="g_normal" type="mesh" mesh="quad" material="mat_normal"/>
+    <geom name="g_plain" type="mesh" mesh="quad" material="mat_plain"/>
+    <geom name="g_blend" type="mesh" mesh="quad" material="mat_plain"
+          rgba="1 1 1 0.5"/>
+  </worldbody>
+</mujoco>
+"""
+
+
 @pytest.fixture
 def hfield_model():
   return mujoco.MjModel.from_xml_string(_HFIELD_XML)
 
 
+@pytest.fixture
+def textured_model():
+  return mujoco.MjModel.from_xml_string(_TEXTURED_XML)
+
+
 @pytest.fixture
 def cubemap_model():
   return mujoco.MjModel.from_xml_string(_CUBEMAP_XML)

tests/test_conversions.py

@@ -2,13 +2,19 @@
 import numpy as np
 import pytest
 import trimesh
+import trimesh.visual
+import trimesh.visual.material
 from mujoco import mjtGeom
+from PIL import Image
 
 from mjviser.conversions import (
   _create_heightfield_mesh,
   _create_shape_mesh,
   _cubemap_vertex_colors,
   _extract_mesh_data,
+  _extract_texture_image,
+  _get_texture_id,
+  _has_alpha,
   _merge_meshes,
   _resolve_flat_rgba,
   create_primitive_mesh,
@@ -120,6 +126,54 @@ def test_cubemap_mesh_to_trimesh(cubemap_model):
   assert len(mesh.vertices) > 0
 
 
+# 2D textures: normal maps and blending
+
+
+def _pbr_material(mesh):
+  assert isinstance(mesh.visual, trimesh.visual.TextureVisuals)
+  mat = mesh.visual.material
+  assert isinstance(mat, trimesh.visual.material.PBRMaterial)
+  return mat
+
+
+def test_textured_mesh_has_pbr_material(textured_model):
+  assert _pbr_material(mujoco_mesh_to_trimesh(textured_model, 0)) is not None
+
+
+def test_normalmap_applied_when_present(textured_model):
+  # Geom 0 uses a material with a normal layer.
+  assert _pbr_material(mujoco_mesh_to_trimesh(textured_model, 0)).normalTexture
+
+
+def test_no_stray_normalmap_when_absent(textured_model):
+  # Geom 1 has an albedo texture but no normal layer; the normalmap texid is
+  # -1 and must not wrap around to another texture.
+  assert _pbr_material(mujoco_mesh_to_trimesh(textured_model, 1)).normalTexture is None
+
+
+def test_opaque_material_is_not_blended(textured_model):
+  assert _pbr_material(mujoco_mesh_to_trimesh(textured_model, 1)).alphaMode == "OPAQUE"
+
+
+def test_translucent_geom_enables_blending(textured_model):
+  # Geom 2 has rgba alpha 0.5.
+  assert _pbr_material(mujoco_mesh_to_trimesh(textured_model, 2)).alphaMode == "BLEND"
+
+
+def test_has_alpha():
+  assert not _has_alpha(Image.new("RGB", (4, 4), (255, 0, 0)))
+  assert not _has_alpha(Image.new("RGBA", (4, 4), (255, 0, 0, 255)))
+  assert _has_alpha(Image.new("RGBA", (4, 4), (255, 0, 0, 128)))
+
+
+def test_extract_texture_flip(textured_model):
+  # Same texture extracted with and without the vertical flip differ.
+  texid = _get_texture_id(textured_model, textured_model.geom_matid[0])
+  flipped = np.asarray(_extract_texture_image(textured_model, texid, flip=True))
+  unflipped = np.asarray(_extract_texture_image(textured_model, texid, flip=False))
+  assert np.array_equal(unflipped, np.flipud(flipped))
+
+
 # Mesh merging