mjviser 0.0.1__py3-none-any.whl

mjviser/__init__.py

@@ -0,0 +1,4 @@
+"""Web-based MuJoCo viewer powered by viser."""
+
+from mjviser.scene import ViserMujocoScene as ViserMujocoScene
+from mjviser.viewer import Viewer as Viewer

mjviser/__main__.py

@@ -0,0 +1,61 @@
+"""Allow ``python -m mjviser model.xml`` to launch the active viewer."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import mujoco
+
+from mjviser import Viewer
+
+
+def _resolve_path(arg: str) -> Path:
+  """Resolve a model path, falling back to recursive glob search in CWD."""
+  path = Path(arg)
+  if path.is_file():
+    return path
+
+  # If pointing at a directory, search for XMLs inside it.
+  if path.is_dir():
+    matches = sorted(path.glob("**/*.xml"))
+  elif "*" in arg:
+    matches = sorted(Path.cwd().glob(arg))
+  else:
+    matches = sorted(Path.cwd().glob(f"**/*{arg}*.xml"))
+  if not matches:
+    print(f"No XML files matching '{arg}' found in {Path.cwd()}")
+    sys.exit(1)
+  if len(matches) == 1:
+    print(f"Found: {matches[0]}")
+    return matches[0]
+
+  # Multiple matches — let the user pick.
+  print(f"Multiple matches for '{arg}':")
+  for i, m in enumerate(matches, 1):
+    try:
+      label = m.relative_to(Path.cwd())
+    except ValueError:
+      label = m
+    print(f"  [{i}] {label}")
+  try:
+    choice = input("Select [1]: ").strip()
+    idx = int(choice) - 1 if choice else 0
+    return matches[idx]
+  except (ValueError, IndexError, KeyboardInterrupt):
+    sys.exit(1)
+
+
+def main() -> None:
+  if len(sys.argv) < 2:
+    print("Usage: python -m mjviser <model.xml>")
+    sys.exit(1)
+
+  path = _resolve_path(sys.argv[1])
+  model = mujoco.MjModel.from_xml_path(str(path))
+  data = mujoco.MjData(model)
+  Viewer(model, data).run()
+
+
+if __name__ == "__main__":
+  main()

mjviser/conversions.py

@@ -0,0 +1,474 @@
+"""Convert MuJoCo mesh data to trimesh format with texture support."""
+
+import mujoco
+import numpy as np
+import trimesh
+import trimesh.visual
+import trimesh.visual.material
+import viser.transforms as vtf
+from mujoco import mj_id2name, mjtGeom, mjtObj
+from PIL import Image
+
+# ------------------------------------------------------------------
+# Color / texture helpers
+# ------------------------------------------------------------------
+
+
+def _get_texture_id(mj_model: mujoco.MjModel, matid: int) -> int:
+  """Return the RGB or RGBA texture ID for a material, or -1."""
+  texid = int(mj_model.mat_texid[matid, int(mujoco.mjtTextureRole.mjTEXROLE_RGB)])
+  if texid < 0:
+    texid = int(mj_model.mat_texid[matid, int(mujoco.mjtTextureRole.mjTEXROLE_RGBA)])
+  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."""
+  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")
+  elif nc in (3, 4):
+    arr = np.flipud(data.reshape(h, w, nc))
+    return Image.fromarray(arr.astype(np.uint8))
+  return None
+
+
+def _resolve_flat_rgba(mj_model: mujoco.MjModel, geom_idx: int) -> np.ndarray:
+  """Resolve the flat RGBA color for a geom as uint8.
+
+  Priority: material rgba > geom rgba.
+  """
+  matid = mj_model.geom_matid[geom_idx]
+  if matid >= 0 and matid < mj_model.nmat:
+    rgba = mj_model.mat_rgba[matid]
+  else:
+    rgba = mj_model.geom_rgba[geom_idx]
+  return (np.clip(rgba, 0, 1) * 255).astype(np.uint8)
+
+
+def _apply_flat_color(mesh: trimesh.Trimesh, rgba_uint8: np.ndarray) -> None:
+  """Apply a uniform RGBA color to all vertices."""
+  mesh.visual = trimesh.visual.ColorVisuals(
+    vertex_colors=np.tile(rgba_uint8, (len(mesh.vertices), 1))
+  )
+
+
+def _cubemap_vertex_colors(
+  mj_model: mujoco.MjModel,
+  matid: int,
+  vertices: np.ndarray,
+  faces: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None:
+  """Project a cube map texture onto mesh triangles by normal direction.
+
+  For each triangle, finds the cube map face most aligned with its
+  normal and assigns that face's average color. If the best face is
+  empty, falls back to the nearest non-empty face (handles chamfered
+  edges).
+
+  Returns (new_vertices, new_faces, vertex_colors) with duplicated
+  vertices so each triangle can have its own color, or None if the
+  material has no cube map texture.
+  """
+  texid = _get_texture_id(mj_model, matid)
+  if texid < 0:
+    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)
+
+  # Average color per cube face. Empty faces stay at [0,0,0,0].
+  face_colors = np.zeros((6, 4), dtype=np.uint8)
+  has_color = np.zeros(6, dtype=bool)
+  for i in range(6):
+    mask = data[i, :, :, : min(nc, 3)].sum(axis=2) > 0
+    if mask.any():
+      avg = data[i][mask].mean(axis=0).astype(np.uint8)
+      face_colors[i, :nc] = avg[:nc]
+      if nc < 4:
+        face_colors[i, 3] = 255
+      has_color[i] = True
+
+  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
+  ranked = np.argsort(-dots, axis=1)
+  nf = len(faces)
+  valid = has_color[ranked]  # (nf, 6) bool
+  first_valid = valid.argmax(axis=1)  # First True per row.
+  best_face = ranked[np.arange(nf), first_valid]
+  tri_colors = face_colors[best_face]
+  # Zero out rows where no face has color.
+  tri_colors[~valid.any(axis=1)] = 0
+
+  # Duplicate vertices so each triangle gets its own color.
+  new_verts = vertices[faces.ravel()]
+  new_faces = np.arange(nf * 3).reshape(-1, 3)
+  vert_colors = np.repeat(tri_colors, 3, axis=0)
+
+  return new_verts, new_faces, vert_colors
+
+
+# ------------------------------------------------------------------
+# Mesh extraction
+# ------------------------------------------------------------------
+
+
+def _extract_mesh_data(
+  mj_model: mujoco.MjModel, geom_idx: int
+) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
+  """Extract vertices, faces, and optional UVs for a mesh geom.
+
+  Returns (vertices, faces, uvs). UVs are None when the mesh has no
+  texture coordinates. When UVs are present, vertices and faces are
+  duplicated so each face-vertex gets its own UV.
+  """
+  mesh_id = mj_model.geom_dataid[geom_idx]
+  vert_start = int(mj_model.mesh_vertadr[mesh_id])
+  vert_count = int(mj_model.mesh_vertnum[mesh_id])
+  face_start = int(mj_model.mesh_faceadr[mesh_id])
+  face_count = int(mj_model.mesh_facenum[mesh_id])
+
+  vertices = mj_model.mesh_vert[vert_start : vert_start + vert_count]
+  faces = mj_model.mesh_face[face_start : face_start + face_count]
+
+  texcoord_num = mj_model.mesh_texcoordnum[mesh_id]
+  if texcoord_num == 0:
+    return vertices, faces, None
+
+  texcoord_adr = mj_model.mesh_texcoordadr[mesh_id]
+  texcoords = mj_model.mesh_texcoord[texcoord_adr : texcoord_adr + texcoord_num]
+  face_tc_idx = mj_model.mesh_facetexcoord[face_start : face_start + face_count]
+
+  # Duplicate vertices so each face-vertex gets its own UV.
+  new_verts = vertices[faces.flatten()]
+  new_uvs = texcoords[face_tc_idx.flatten()]
+  new_faces = np.arange(face_count * 3).reshape(-1, 3)
+
+  return new_verts, new_faces, new_uvs
+
+
+def mujoco_mesh_to_trimesh(mj_model: mujoco.MjModel, geom_idx: int) -> trimesh.Trimesh:
+  """Convert a MuJoCo mesh geometry to a trimesh with visual appearance.
+
+  Color resolution order:
+    1. 2D texture with UVs (PBR material with baseColorTexture)
+    2. Cube map texture projected by triangle normals
+    3. Flat material color (mat_rgba)
+    4. Flat geom color (geom_rgba)
+  """
+  vertices, faces, uvs = _extract_mesh_data(mj_model, geom_idx)
+  mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
+
+  matid = mj_model.geom_matid[geom_idx]
+
+  # 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
+
+  # Path 2: no UVs, try cube map projection.
+  if uvs is None and matid >= 0:
+    result = _cubemap_vertex_colors(mj_model, matid, vertices, faces)
+    if result is not None:
+      new_verts, new_faces, vert_colors = result
+      mesh = trimesh.Trimesh(vertices=new_verts, faces=new_faces, process=False)
+      mesh.visual = trimesh.visual.ColorVisuals(vertex_colors=vert_colors)
+      return mesh
+
+  # Path 3/4: flat color from material or geom.
+  _apply_flat_color(mesh, _resolve_flat_rgba(mj_model, geom_idx))
+  return mesh
+
+
+def _create_shape_mesh(geom_type: int, size: np.ndarray) -> trimesh.Trimesh:
+  """Create an uncolored mesh for a standard shape type."""
+  if geom_type == mjtGeom.mjGEOM_SPHERE:
+    return trimesh.creation.icosphere(radius=size[0], subdivisions=2)
+  elif geom_type == mjtGeom.mjGEOM_BOX:
+    return trimesh.creation.box(extents=2.0 * size)
+  elif geom_type == mjtGeom.mjGEOM_CAPSULE:
+    return trimesh.creation.capsule(radius=size[0], height=2.0 * size[1])
+  elif geom_type == mjtGeom.mjGEOM_CYLINDER:
+    return trimesh.creation.cylinder(radius=size[0], height=2.0 * size[1])
+  elif geom_type == mjtGeom.mjGEOM_ELLIPSOID:
+    mesh = trimesh.creation.icosphere(subdivisions=3, radius=1.0)
+    mesh.apply_scale(size)
+    return mesh
+  raise ValueError(f"Unsupported shape type: {geom_type}")
+
+
+def create_primitive_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Trimesh:
+  """Create a mesh for primitive geom types.
+
+  Supports sphere, box, capsule, cylinder, plane, ellipsoid, and
+  heightfield.
+  """
+  geom_type = mj_model.geom_type[geom_id]
+
+  if geom_type == mjtGeom.mjGEOM_HFIELD:
+    return _create_heightfield_mesh(mj_model, geom_id)
+
+  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
+    plane_y = 2.0 * size[1] if size[1] > 0 else 20.0
+    mesh = trimesh.creation.box((plane_x, plane_y, 0.001))
+  else:
+    mesh = _create_shape_mesh(geom_type, mj_model.geom_size[geom_id])
+
+  _apply_flat_color(mesh, _resolve_flat_rgba(mj_model, geom_id))
+  return mesh
+
+
+def _create_heightfield_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Trimesh:
+  """Create a colored heightfield mesh."""
+  hfield_id = mj_model.geom_dataid[geom_id]
+  nrow = mj_model.hfield_nrow[hfield_id]
+  ncol = mj_model.hfield_ncol[hfield_id]
+  sx, sy, sz, _base = mj_model.hfield_size[hfield_id]
+
+  offset = 0
+  for k in range(hfield_id):
+    offset += mj_model.hfield_nrow[k] * mj_model.hfield_ncol[k]
+  hfield = mj_model.hfield_data[offset : offset + nrow * ncol].reshape(nrow, ncol)
+
+  x = np.linspace(-sx, sx, ncol)
+  y = np.linspace(-sy, sy, nrow)
+  xx, yy = np.meshgrid(x, y)
+  zz = hfield * sz
+
+  vertices = np.column_stack((xx.ravel(), yy.ravel(), zz.ravel()))
+
+  ri, ci = np.mgrid[: nrow - 1, : ncol - 1]
+  i0 = (ri * ncol + ci).ravel()
+  faces = np.column_stack(
+    [i0, i0 + 1, i0 + ncol + 1, i0, i0 + ncol + 1, i0 + ncol]
+  ).reshape(-1, 3)
+  mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
+
+  # Color by height using HSV.
+  zz_min, zz_max = zz.min(), zz.max()
+  normalized = (
+    (zz - zz_min) / (zz_max - zz_min) if zz_max > zz_min else np.full_like(zz, 0.5)
+  )
+
+  hue = 0.5 - normalized * 0.45
+  saturation = 0.6 - normalized * 0.2
+  value = 0.4 + normalized * 0.3
+
+  c = value * saturation
+  x = c * (1 - np.abs((hue * 6) % 2 - 1))
+  m = value - c
+
+  hue_sector = (hue * 6).astype(int) % 6
+  r = np.zeros_like(hue)
+  g = np.zeros_like(hue)
+  b = np.zeros_like(hue)
+
+  for sector, rc, gc, bc in [
+    (0, c, x, 0),
+    (1, x, c, 0),
+    (2, 0, c, x),
+    (3, 0, x, c),
+    (4, x, 0, c),
+    (5, c, 0, x),
+  ]:
+    mask = hue_sector == sector
+    r[mask] = rc[mask] if isinstance(rc, np.ndarray) else rc
+    g[mask] = gc[mask] if isinstance(gc, np.ndarray) else gc
+    b[mask] = bc[mask] if isinstance(bc, np.ndarray) else bc
+
+  r += m
+  g += m
+  b += m
+
+  vertex_colors = np.column_stack(
+    [
+      (np.clip(r, 0, 1) * 255).astype(np.uint8).ravel(),
+      (np.clip(g, 0, 1) * 255).astype(np.uint8).ravel(),
+      (np.clip(b, 0, 1) * 255).astype(np.uint8).ravel(),
+      np.full(vertices.shape[0], 255, dtype=np.uint8),
+    ]
+  )
+
+  mesh.visual = trimesh.visual.ColorVisuals(mesh=mesh, vertex_colors=vertex_colors)
+  return mesh
+
+
+# ------------------------------------------------------------------
+# Texture / visual utilities
+# ------------------------------------------------------------------
+
+
+def get_geom_texture_id(mj_model: mujoco.MjModel, geom_idx: int) -> int:
+  """Return the texture ID the trimesh conversion will use, or -1.
+
+  Returns -1 when the geom is a primitive, the material has no
+  texture, or the mesh has no UV coordinates.
+  """
+  if mj_model.geom_type[geom_idx] != mjtGeom.mjGEOM_MESH:
+    return -1
+
+  matid = mj_model.geom_matid[geom_idx]
+  if matid < 0 or matid >= mj_model.nmat:
+    return -1
+
+  texid = _get_texture_id(mj_model, matid)
+  if texid < 0:
+    return -1
+
+  mesh_id = mj_model.geom_dataid[geom_idx]
+  if mj_model.mesh_texcoordnum[mesh_id] <= 0:
+    return -1
+
+  return texid
+
+
+def group_geoms_by_visual_compat(
+  mj_model: mujoco.MjModel, geom_ids: list[int]
+) -> list[list[int]]:
+  """Partition geom IDs into groups that can be safely merged.
+
+  Geoms sharing the same texture ID are grouped together. All
+  untextured geoms form a single group.
+  """
+  groups: dict[int, list[int]] = {}
+  for gid in geom_ids:
+    tex_id = get_geom_texture_id(mj_model, gid)
+    groups.setdefault(tex_id, []).append(gid)
+  return list(groups.values())
+
+
+# ------------------------------------------------------------------
+# Merge / transform utilities
+# ------------------------------------------------------------------
+
+
+def _merge_meshes(
+  meshes: list[trimesh.Trimesh],
+  positions: list[np.ndarray],
+  quats: list[np.ndarray],
+) -> trimesh.Trimesh:
+  """Transform and merge meshes given positions and wxyz quaternions."""
+  for mesh, pos, quat in zip(meshes, positions, quats, strict=True):
+    transform = np.eye(4)
+    transform[:3, :3] = vtf.SO3(quat).as_matrix()
+    transform[:3, 3] = pos
+    mesh.apply_transform(transform)
+
+  result = meshes[0] if len(meshes) == 1 else trimesh.util.concatenate(meshes)
+  result.merge_vertices()
+  return result
+
+
+def merge_geoms(mj_model: mujoco.MjModel, geom_ids: list[int]) -> trimesh.Trimesh:
+  """Merge multiple geoms into a single trimesh in local body space."""
+  meshes = []
+  for geom_id in geom_ids:
+    if mj_model.geom_type[geom_id] == mjtGeom.mjGEOM_MESH:
+      meshes.append(mujoco_mesh_to_trimesh(mj_model, geom_id))
+    else:
+      meshes.append(create_primitive_mesh(mj_model, geom_id))
+
+  return _merge_meshes(
+    meshes,
+    [mj_model.geom_pos[gid] for gid in geom_ids],
+    [mj_model.geom_quat[gid] for gid in geom_ids],
+  )
+
+
+def rotation_matrix_from_vectors(
+  from_vec: np.ndarray, to_vec: np.ndarray
+) -> np.ndarray:
+  """3x3 rotation matrix that rotates from_vec to to_vec (Rodrigues)."""
+  from_vec = from_vec / np.linalg.norm(from_vec)
+  to_vec = to_vec / np.linalg.norm(to_vec)
+
+  if np.allclose(from_vec, to_vec):
+    return np.eye(3)
+  if np.allclose(from_vec, -to_vec):
+    return np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]])
+
+  v = np.cross(from_vec, to_vec)
+  s = np.linalg.norm(v)
+  c = np.dot(from_vec, to_vec)
+  vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
+  return np.eye(3) + vx + vx @ vx * ((1 - c) / (s * s))
+
+
+def is_fixed_body(mj_model: mujoco.MjModel, body_id: int) -> bool:
+  """Check if a body is fixed (welded to world, not attached to mocap)."""
+  is_weld = mj_model.body_weldid[body_id] == 0
+  root_id = mj_model.body_rootid[body_id]
+  return bool(is_weld and mj_model.body_mocapid[root_id] < 0)
+
+
+def get_body_name(mj_model: mujoco.MjModel, body_id: int) -> str:
+  """Body name with fallback to ``body_{id}``."""
+  name = mj_id2name(mj_model, mjtObj.mjOBJ_BODY, body_id)
+  return name if name else f"body_{body_id}"
+
+
+def create_site_mesh(mj_model: mujoco.MjModel, site_id: int) -> trimesh.Trimesh:
+  """Create a mesh for a single site."""
+  rgba = mj_model.site_rgba[site_id].copy()
+  if np.all(rgba == 0):
+    rgba = np.array([0.5, 0.5, 0.5, 1.0])
+  rgba_uint8 = (np.clip(rgba, 0, 1) * 255).astype(np.uint8)
+
+  mesh = _create_shape_mesh(mj_model.site_type[site_id], mj_model.site_size[site_id])
+  _apply_flat_color(mesh, rgba_uint8)
+  return mesh
+
+
+def merge_sites(mj_model: mujoco.MjModel, site_ids: list[int]) -> trimesh.Trimesh:
+  """Merge multiple sites into a single trimesh in local body space."""
+  return _merge_meshes(
+    [create_site_mesh(mj_model, sid) for sid in site_ids],
+    [mj_model.site_pos[sid] for sid in site_ids],
+    [mj_model.site_quat[sid] for sid in site_ids],
+  )
+
+
+def get_site_name(mj_model: mujoco.MjModel, site_id: int) -> str:
+  """Site name with fallback to ``site_{id}``."""
+  name = mj_id2name(mj_model, mjtObj.mjOBJ_SITE, site_id)
+  return name if name else f"site_{site_id}"