Render cube map textures on box geoms

src/mjviser/conversions.py

@@ -34,6 +34,24 @@ def _has_alpha(image: Image.Image) -> bool:
   return bool(np.asarray(image.getchannel("A")).min() < 255)
 
 
+def _is_cubemap_texture(mj_model: mujoco.MjModel, texid: int) -> bool:
+  """Return True if texid is a cube map stored as 6 stacked square faces."""
+  if int(mj_model.tex_type[texid]) != int(mujoco.mjtTexture.mjTEXTURE_CUBE):
+    return False
+  w = int(mj_model.tex_width[texid])
+  h = int(mj_model.tex_height[texid])
+  nc = int(mj_model.tex_nchannel[texid])
+  return h == w * 6 and nc in (1, 3, 4)
+
+
+# Cube map face order with the outward scene axis each face sits on:
+# +X (right), -X (left), +Y (up), -Y (down), +Z (front), -Z (back).
+_CUBEMAP_AXES: np.ndarray = np.array(
+  [[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]],
+  dtype=np.float64,
+)
+
+
 def _extract_texture_image(
   mj_model: mujoco.MjModel, texid: int, flip: bool = True
 ) -> Image.Image | None:
@@ -102,15 +120,12 @@ def _cubemap_vertex_colors(
   material has no cube map texture.
   """
   texid = _get_texture_id(mj_model, matid)
-  if texid < 0:
+  if texid < 0 or not _is_cubemap_texture(mj_model, texid):
     return None
 
   w = mj_model.tex_width[texid]
   h = mj_model.tex_height[texid]
   nc = mj_model.tex_nchannel[texid]
-  if int(mj_model.tex_type[texid]) != 1 or h != w * 6:
-    return None
-
   adr = mj_model.tex_adr[texid]
   data = mj_model.tex_data[adr : adr + w * h * nc].reshape(6, w, w, nc)
 
@@ -129,19 +144,13 @@ def _cubemap_vertex_colors(
   if not has_color.any():
     return None
 
-  # Cube map axes: +X, -X, +Y, -Y, +Z, -Z.
-  axes = np.array(
-    [[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]],
-    dtype=np.float64,
-  )
-
   # Per-triangle normals.
   v0, v1, v2 = vertices[faces[:, 0]], vertices[faces[:, 1]], vertices[faces[:, 2]]
   normals = np.cross(v1 - v0, v2 - v0)
   normals /= np.maximum(np.linalg.norm(normals, axis=1, keepdims=True), 1e-10)
 
   # For each triangle, pick the best aligned face that has color.
-  dots = normals @ axes.T
+  dots = normals @ _CUBEMAP_AXES.T
   ranked = np.argsort(-dots, axis=1)
   nf = len(faces)
   valid = has_color[ranked]  # (nf, 6) bool
@@ -283,6 +292,12 @@ def create_primitive_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Tri
   if geom_type == mjtGeom.mjGEOM_HFIELD:
     return _create_heightfield_mesh(mj_model, geom_id)
 
+  # Box primitives with a cube map texture get per-face textured rendering.
+  if geom_type == mjtGeom.mjGEOM_BOX:
+    textured = _create_cubemap_box_mesh(mj_model, geom_id)
+    if textured is not None:
+      return textured
+
   if geom_type == mjtGeom.mjGEOM_PLANE:
     size = mj_model.geom_size[geom_id]
     plane_x = 2.0 * size[0] if size[0] > 0 else 20.0
@@ -295,6 +310,114 @@ def create_primitive_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Tri
   return mesh
 
 
+# Per cube face, the (image-right, image-up) directions in scene coordinates as
+# seen by a viewer looking at the face from outside the cube. Indexed by the
+# face order in _CUBEMAP_AXES. Matches MuJoCo's cube map convention.
+_CUBEMAP_FACE_BASIS: tuple[tuple[np.ndarray, np.ndarray], ...] = (
+  (np.array([0.0, -1.0, 0.0]), np.array([0.0, 0.0, 1.0])),  # +X
+  (np.array([0.0, 1.0, 0.0]), np.array([0.0, 0.0, 1.0])),  # -X
+  (np.array([1.0, 0.0, 0.0]), np.array([0.0, 0.0, -1.0])),  # +Y
+  (np.array([1.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0])),  # -Y
+  (np.array([1.0, 0.0, 0.0]), np.array([0.0, 1.0, 0.0])),  # +Z
+  (np.array([-1.0, 0.0, 0.0]), np.array([0.0, 1.0, 0.0])),  # -Z
+)
+
+
+def _extract_cubemap_atlas(mj_model: mujoco.MjModel, texid: int) -> Image.Image | None:
+  """Build a vertical-strip PIL atlas from a cube map texture.
+
+  Returns an image of size (w, 6*w) with face i pasted at rows
+  [i*w, (i+1)*w) in PIL top-down coordinates, or None if the texture isn't a
+  supported cube map.
+  """
+  if not _is_cubemap_texture(mj_model, texid):
+    return None
+  w = int(mj_model.tex_width[texid])
+  h = int(mj_model.tex_height[texid])
+  nc = int(mj_model.tex_nchannel[texid])
+  adr = int(mj_model.tex_adr[texid])
+  data = mj_model.tex_data[adr : adr + w * h * nc].reshape(6, w, w, nc)
+  mode = {1: "L", 3: "RGB", 4: "RGBA"}[nc]
+
+  atlas = Image.new(mode, (w, 6 * w))
+  for i in range(6):
+    # MuJoCo stores rows bottom-up; flip to PIL top-down so the face
+    # image appears right-side-up in the atlas.
+    arr = np.flipud(data[i]).astype(np.uint8)
+    if nc == 1:
+      arr = arr.reshape(w, w)
+    atlas.paste(Image.fromarray(arr, mode=mode), (0, i * w))
+  return atlas
+
+
+def _create_cubemap_box_mesh(
+  mj_model: mujoco.MjModel, geom_id: int
+) -> trimesh.Trimesh | None:
+  """Build a 6-quad textured box mesh from the geom's cube map material.
+
+  Returns None when the geom has no material, no cube map texture, or
+  the texture has an unsupported format. Each cube face gets its own
+  slice of a vertical-strip atlas so the per-face images render with
+  correct orientation.
+  """
+  matid = int(mj_model.geom_matid[geom_id])
+  if matid < 0 or matid >= mj_model.nmat:
+    return None
+  texid = _get_texture_id(mj_model, matid)
+  if texid < 0:
+    return None
+  atlas = _extract_cubemap_atlas(mj_model, texid)
+  if atlas is None:
+    return None
+
+  sx, sy, sz = mj_model.geom_size[geom_id]
+  centers = _CUBEMAP_AXES * np.array([sx, sx, sy, sy, sz, sz])[:, None]
+  half_extents = [(sy, sz), (sy, sz), (sx, sz), (sx, sz), (sx, sy), (sx, sy)]
+
+  verts = np.zeros((24, 3), dtype=np.float64)
+  uvs = np.zeros((24, 2), dtype=np.float64)
+  faces = np.zeros((12, 3), dtype=np.int64)
+  for fi in range(6):
+    c = centers[fi]
+    r, u = _CUBEMAP_FACE_BASIS[fi]
+    hr, hu = half_extents[fi]
+    # Corners in image-relative order: bottom-left, bottom-right, top-right,
+    # top-left.
+    base = fi * 4
+    verts[base + 0] = c - hr * r - hu * u
+    verts[base + 1] = c + hr * r - hu * u
+    verts[base + 2] = c + hr * r + hu * u
+    verts[base + 3] = c - hr * r + hu * u
+    # Atlas UVs (OpenGL convention: v=0 at bottom of atlas image).
+    v_top = 1.0 - fi / 6.0
+    v_bot = 1.0 - (fi + 1) / 6.0
+    uvs[base + 0] = (0.0, v_bot)
+    uvs[base + 1] = (1.0, v_bot)
+    uvs[base + 2] = (1.0, v_top)
+    uvs[base + 3] = (0.0, v_top)
+    # Pick triangle winding so the face normal points outward. The image
+    # basis (r, u) is chosen so the texture reads correctly when viewed
+    # from outside, which means r × u may point either way relative to
+    # the outward axis depending on the face.
+    if np.dot(np.cross(r, u), _CUBEMAP_AXES[fi]) >= 0:
+      faces[fi * 2 + 0] = (base + 0, base + 1, base + 2)
+      faces[fi * 2 + 1] = (base + 0, base + 2, base + 3)
+    else:
+      faces[fi * 2 + 0] = (base + 0, base + 2, base + 1)
+      faces[fi * 2 + 1] = (base + 0, base + 3, base + 2)
+
+  mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False)
+  rgba = mj_model.mat_rgba[matid]
+  material = trimesh.visual.material.PBRMaterial(
+    baseColorFactor=rgba,
+    baseColorTexture=atlas,
+    metallicFactor=0.0,
+    roughnessFactor=1.0,
+  )
+  mesh.visual = trimesh.visual.TextureVisuals(uv=uvs, material=material)
+  return mesh
+
+
 def _create_heightfield_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Trimesh:
   """Create a heightfield mesh, using the material texture when available."""
   hfield_id = mj_model.geom_dataid[geom_id]
@@ -377,6 +500,10 @@ def get_geom_texture_id(mj_model: mujoco.MjModel, geom_idx: int) -> int:
   if geom_type == mjtGeom.mjGEOM_HFIELD:
     return texid
 
+  # Box primitives use the cube map texture via per-face UV mapping.
+  if geom_type == mjtGeom.mjGEOM_BOX:
+    return texid if _is_cubemap_texture(mj_model, texid) else -1
+
   if geom_type != mjtGeom.mjGEOM_MESH:
     return -1
 

tests/test_conversions.py

@@ -20,6 +20,7 @@
   create_primitive_mesh,
   create_site_mesh,
   get_body_name,
+  get_geom_texture_id,
   group_geoms_by_visual_compat,
   is_fixed_body,
   merge_geoms,
@@ -129,9 +130,13 @@ def test_cubemap_mesh_to_trimesh(cubemap_model):
 # 2D textures: normal maps and blending
 
 
-def _pbr_material(mesh):
+def _texture_visual(mesh):
   assert isinstance(mesh.visual, trimesh.visual.TextureVisuals)
-  mat = mesh.visual.material
+  return mesh.visual
+
+
+def _pbr_material(mesh):
+  mat = _texture_visual(mesh).material
   assert isinstance(mat, trimesh.visual.material.PBRMaterial)
   return mat
 
@@ -174,6 +179,113 @@ def test_extract_texture_flip(textured_model):
   assert np.array_equal(unflipped, np.flipud(flipped))
 
 
+# Cube map textures on box primitives
+
+
+_CUBEMAP_BOX_XML = """
+<mujoco>
+  <asset>
+    <texture name="cubetex" type="cube" builtin="flat" mark="cross"
+             width="32" height="32" rgb1="0.8 0.2 0.2" markrgb="1 1 1"/>
+    <texture name="tex2d" type="2d" builtin="checker"
+             width="32" height="32" rgb1="0.8 0.2 0.2" rgb2="0.2 0.2 0.8"/>
+    <material name="cubemat" texture="cubetex"/>
+    <material name="flatmat" texture="tex2d"/>
+  </asset>
+  <worldbody>
+    <geom name="cube" type="box" size="0.05 0.06 0.07" material="cubemat"/>
+    <geom name="plain" type="box" size="0.05 0.05 0.05" rgba="1 0 0 1"/>
+    <geom name="box2d" type="box" size="0.05 0.05 0.05" material="flatmat"/>
+  </worldbody>
+</mujoco>
+"""
+
+# MuJoCo cube face order with the outward axis each face sits on.
+_CUBE_FACE_AXES = (
+  (0, (1, 0, 0)),
+  (1, (-1, 0, 0)),
+  (2, (0, 1, 0)),
+  (3, (0, -1, 0)),
+  (4, (0, 0, 1)),
+  (5, (0, 0, -1)),
+)
+
+
+def test_cubemap_box_mesh_structure():
+  # A box with a cube map material becomes a textured 6-quad mesh whose
+  # triangles all wind outward (closed, positive-volume box).
+  model = mujoco.MjModel.from_xml_string(_CUBEMAP_BOX_XML)
+  mesh = create_primitive_mesh(model, 0)
+  material = _pbr_material(mesh)
+  assert mesh.vertices.shape == (24, 3)
+  assert mesh.faces.shape == (12, 3)
+  assert _texture_visual(mesh).uv.shape == (24, 2)
+  # 6 square faces stacked into a vertical strip: atlas is (w, 6*w).
+  atlas = material.baseColorTexture
+  assert atlas is not None and atlas.size == (32, 192)
+  assert mesh.volume > 0
+  assert mesh.is_winding_consistent
+
+
+def test_cubemap_box_extents_match_geom_size():
+  # The textured box spans the full geom size on every axis.
+  model = mujoco.MjModel.from_xml_string(_CUBEMAP_BOX_XML)
+  mesh = create_primitive_mesh(model, 0)
+  lo, hi = mesh.bounds
+  np.testing.assert_allclose(hi, [0.05, 0.06, 0.07])
+  np.testing.assert_allclose(lo, [-0.05, -0.06, -0.07])
+
+
+def test_cubemap_box_face_order_and_placement():
+  # Paint each cube face a distinct color, then check each quad sits on the
+  # correct outward axis and samples its own face color from the atlas.
+  model = mujoco.MjModel.from_xml_string(_CUBEMAP_BOX_XML)
+  colors = np.array(
+    [
+      [200, 0, 0],
+      [0, 200, 0],
+      [0, 0, 200],
+      [200, 200, 0],
+      [200, 0, 200],
+      [0, 200, 200],
+    ],
+    dtype=np.uint8,
+  )
+  w = int(model.tex_width[0])
+  adr = int(model.tex_adr[0])
+  for i, color in enumerate(colors):
+    start = adr + i * w * w * 3
+    model.tex_data[start : start + w * w * 3] = np.tile(color, w * w)
+
+  mesh = create_primitive_mesh(model, 0)
+  atlas = np.asarray(_pbr_material(mesh).baseColorTexture)
+  uv = _texture_visual(mesh).uv
+  for fi, axis in _CUBE_FACE_AXES:
+    quad = slice(fi * 4, fi * 4 + 4)
+    centroid = mesh.vertices[quad].mean(axis=0)
+    direction = centroid / np.linalg.norm(centroid)
+    np.testing.assert_allclose(direction, axis, atol=1e-6)
+    u, v = uv[quad].mean(axis=0)
+    px = min(int(u * atlas.shape[1]), atlas.shape[1] - 1)
+    py = min(int((1.0 - v) * atlas.shape[0]), atlas.shape[0] - 1)
+    np.testing.assert_allclose(atlas[py, px], colors[fi], atol=2)
+
+
+def test_box_without_cube_map_is_flat_colored():
+  # Plain boxes and boxes with a 2D (non-cube) texture keep the flat fallback.
+  model = mujoco.MjModel.from_xml_string(_CUBEMAP_BOX_XML)
+  assert isinstance(create_primitive_mesh(model, 1).visual, trimesh.visual.ColorVisuals)
+  assert isinstance(create_primitive_mesh(model, 2).visual, trimesh.visual.ColorVisuals)
+
+
+def test_get_geom_texture_id_cube_map_box():
+  # Boxes group by cube map texture; plain and 2D-textured boxes do not.
+  model = mujoco.MjModel.from_xml_string(_CUBEMAP_BOX_XML)
+  assert get_geom_texture_id(model, 0) == 0
+  assert get_geom_texture_id(model, 1) == -1
+  assert get_geom_texture_id(model, 2) == -1
+
+
 # Mesh merging