diff --git a/src/mjviser/conversions.py b/src/mjviser/conversions.py index 1a9d02c..4ebbf72 100644 --- a/src/mjviser/conversions.py +++ b/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: diff --git a/tests/conftest.py b/tests/conftest.py index 0dac624..d5e8666 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -83,11 +83,41 @@ def simple_data(simple_model): """ +_TEXTURED_XML = """ + + + + + + + + + + + + + + + + + +""" + + @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) diff --git a/tests/test_conversions.py b/tests/test_conversions.py index b65156f..d1f3a43 100644 --- a/tests/test_conversions.py +++ b/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