Add mesh conversion utilities.

PiperOrigin-RevId: 573343028
Change-Id: I4b258f91a11d1994659b0129de5ddb2077689f01

python/build_requirements.txt

@@ -2,6 +2,10 @@ absl-py==2.0.0 \
     --hash=sha256:9a28abb62774ae4e8edbe2dd4c49ffcd45a6a848952a5eccc6a49f3f0fc1e2f3
 auditwheel==5.4.0; platform_system == 'Linux' \
     --hash=sha256:d8410a17523427ba3f7b60c9701d23de28b8f94fa5dab732aa6c30d160df8127
+etils[epath]==1.3.0; python_version == '3.8' \
+    --hash=sha256:809a92ff72f12149441492cf4d9a26b56a4741dffb4dfb9c4c7b7afe055c2d28
+etils[epath]==1.5.1; python_version >= '3.9' \
+    --hash=sha256:2c1bfa2817eb4881cb509097f1e65ac6160126ba74ec47b3bb47ee678628d8c8
 glfw==2.6.2 \
     --hash=sha256:c2dcf2395d99ff2506428213bee305bb9ba024043d1f574216e61e6f5df808e9 \
     --hash=sha256:c385c9976133aed57ff4f0ff5210276844cefb4d8a7bf61bbcc1caf10385744b \
@@ -50,6 +54,16 @@ wheel==0.41.2 \
 pyelftools==0.30; platform_system == 'Linux' \
     --hash=sha256:544c3440eddb9a0dce70b6611de0b28163d71def759d2ed57a0d00118fc5da86
 
+# Transitive dependencies of etils[epath]
+fsspec==2023.9.2 \
+    --hash=sha256:603dbc52c75b84da501b9b2ec8c11e1f61c25984c4a0dda1f129ef391fbfc9b4
+importlib-resources==6.1.0 \
+    --hash=sha256:aa50258bbfa56d4e33fbd8aa3ef48ded10d1735f11532b8df95388cc6bdb7e83
+typing_extensions==4.8.0 \
+    --hash=sha256:8f92fc8806f9a6b641eaa5318da32b44d401efaac0f6678c9bc448ba3605faa0
+zipp==3.17.0 \
+    --hash=sha256:0e923e726174922dce09c53c59ad483ff7bbb8e572e00c7f7c46b88556409f31
+
 # Transitive dependencies of pytest
 attrs==23.1.0; platform_system == 'Windows' \
     --hash=sha256:1f28b4522cdc2fb4256ac1a020c78acf9cba2c6b461ccd2c126f3aa8e8335d04

python/mujoco/msh2obj.py

@@ -0,0 +1,111 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""CLI for converting legacy MSH files to Wavefront OBJ files.
+
+Usage:
+  python -m mujoco.msh2obj -i <msh_file> -o <obj_file>
+"""
+
+import argparse
+import dataclasses
+import io
+import pathlib
+
+import numpy as np
+
+
+@dataclasses.dataclass(frozen=True)
+class Msh:
+  """MuJoCo legacy binary msh file."""
+
+  vertex_positions: np.ndarray
+  vertex_normals: np.ndarray
+  vertex_texcoords: np.ndarray
+  face_vertex_indices: np.ndarray
+
+  @staticmethod
+  def create(file: pathlib.Path) -> "Msh":
+    """Create a Msh object from a .msh file."""
+    if not file.exists():
+      raise FileNotFoundError(f"{file} does not exist.")
+
+    with open(file, "rb") as f:
+      nvertex = np.fromfile(f, dtype=np.int32, count=1)[0]
+      nnormal = np.fromfile(f, dtype=np.int32, count=1)[0]
+      ntexcoord = np.fromfile(f, dtype=np.int32, count=1)[0]
+      nface = np.fromfile(f, dtype=np.int32, count=1)[0]
+      vertex_positions = np.fromfile(f, dtype=np.float32, count=3 * nvertex)
+      vertex_normals = np.fromfile(f, dtype=np.float32, count=3 * nnormal)
+      vertex_texcoords = np.fromfile(f, dtype=np.float32, count=2 * ntexcoord)
+      face_vertex_indices = np.fromfile(f, dtype=np.int32, count=3 * nface)
+
+    if vertex_positions.size != 3 * nvertex:
+      raise ValueError(
+          f"Invalid number of vertices: {vertex_positions.size} != 3*{nvertex}."
+      )
+    if vertex_normals.size != 3 * nnormal:
+      raise ValueError(
+          f"Invalid number of normals: {vertex_normals.size} != 3*{nnormal}."
+      )
+    if vertex_texcoords.size != 2 * ntexcoord:
+      raise ValueError(
+          f"Invalid number of texcoords: {vertex_texcoords.size} != "
+          "2*{ntexcoord}."
+      )
+    if face_vertex_indices.size != 3 * nface:
+      raise ValueError(
+          f"Invalid number of faces: {face_vertex_indices.size} != 3*{nface}."
+      )
+
+    vertex_positions = vertex_positions.reshape(-1, 3)
+    vertex_normals = vertex_normals.reshape(-1, 3)
+    face_vertex_indices = face_vertex_indices.reshape(-1, 3)
+
+    # Undo vertical flip done by MuJoCo's OBJ loader.
+    vertex_texcoords = vertex_texcoords.reshape(-1, 2)
+    vertex_texcoords[:, 1] = 1 - vertex_texcoords[:, 1]
+
+    return Msh(
+        vertex_positions=vertex_positions,
+        vertex_normals=vertex_normals,
+        vertex_texcoords=vertex_texcoords,
+        face_vertex_indices=face_vertex_indices,
+    )
+
+
+def msh_to_obj(msh_file: pathlib.Path) -> str:
+  """Convert a legacy .msh file to the .obj format."""
+  msh = Msh.create(msh_file)
+
+  out = io.StringIO()
+  for x, y, z in msh.vertex_positions:
+    out.write(f"v {x} {y} {z}\n")
+  for x, y, z in msh.vertex_normals:
+    out.write(f"vn {x} {y} {z}\n")
+  for u, v in msh.vertex_texcoords:
+    out.write(f"vt {u} {v}\n")
+  for i, j, k in msh.face_vertex_indices:
+    out.write(f"f {i+1}/{i+1}/{i+1} {j+1}/{j+1}/{j+1} {k+1}/{k+1}/{k+1}\n")
+
+  return out.getvalue()
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument("-i", "--input", type=str, help="Path to the msh file.")
+  parser.add_argument("-o", "--output", type=str, help="Path to the obj file.")
+  args = parser.parse_args()
+  with open(pathlib.Path(args.output), "w") as f:
+    f.write(msh_to_obj(pathlib.Path(args.input)))

python/mujoco/msh2obj_test.py

@@ -0,0 +1,77 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for msh2obj.py."""
+
+from absl.testing import absltest
+from etils import epath
+import mujoco
+from mujoco import msh2obj
+import numpy as np
+
+
+_MESH_FIELDS = (
+    "mesh_vertadr",
+    "mesh_vertnum",
+    "mesh_faceadr",
+    "mesh_facenum",
+    "mesh_bvhadr",
+    "mesh_bvhnum",
+    "mesh_normaladr",
+    "mesh_normalnum",
+    "mesh_texcoordadr",
+    "mesh_texcoordnum",
+    "mesh_graphadr",
+    "mesh_vert",
+    "mesh_normal",
+    "mesh_face",
+    "mesh_facenormal",
+    "mesh_facetexcoord",
+    "mesh_graph",
+    "mesh_texcoord",
+)
+
+_XML = """
+<mujoco>
+  <asset>
+    <mesh name="abdomen_1_body" file="abdomen_1_body.obj"/>
+  </asset>
+</mujoco>
+"""
+
+
+class MshTest(absltest.TestCase):
+
+  def test_obj_model_matches_msh_model(self) -> None:
+    test_path = epath.resource_path("mujoco") / "testdata"
+
+    msh_xml = test_path / "msh.xml"
+    msh_model = mujoco.MjModel.from_xml_path(msh_xml.as_posix())
+
+    msh_path = test_path / "abdomen_1_body.msh"
+    obj = msh2obj.msh_to_obj(msh_path)
+
+    obj_model = mujoco.MjModel.from_xml_string(
+        _XML, {"abdomen_1_body.obj": obj.encode()})
+
+    for field in _MESH_FIELDS:
+      np.testing.assert_allclose(
+          getattr(msh_model, field),
+          getattr(obj_model, field),
+          atol=1e-6,
+          err_msg=f"Field {field} does not match between msh and obj models.",
+      )
+
+if __name__ == "__main__":
+  absltest.main()

python/mujoco/testdata/msh.xml

@@ -0,0 +1,5 @@
+<mujoco>
+  <asset>
+    <mesh name="abdomen_1_body" file="abdomen_1_body.msh"/>
+  </asset>
+</mujoco>

python/pyproject.toml

@@ -27,6 +27,7 @@ classifiers = [
 ]
 dependencies = [
     "absl-py",
+    "etils[epath]",
     "glfw",
     "numpy",
     "pyopengl",
@@ -55,4 +56,6 @@ mujoco = [
     "libmujoco*.so.*",
     "mujoco.dll",
     "include/mujoco/*.h",
+    "testdata/*.xml",
+    "testdata/*.msh",
 ]