removed processing from TorchMeshContainer

Author: mhubii

roboreg/differentiable/structs.py

@@ -2,10 +2,10 @@
 from collections import OrderedDict
 from typing import Dict, List, Optional, Tuple, Union
 
-import fast_simplification
 import numpy as np
 import torch
-import trimesh
+
+from roboreg.io import Mesh
 
 
 class TorchMeshContainer:
@@ -17,7 +17,7 @@ class TorchMeshContainer:
     """
 
     __slots__ = [
-        "_mesh_names",
+        "_names",
         "_vertices",  # tensor of shape (B, N, 4) -> homogeneous coordinates
         "_faces",  # tensor of shape (B, N, 3)
         "_per_mesh_vertex_count",
@@ -32,12 +32,11 @@ class TorchMeshContainer:
 
     def __init__(
         self,
-        mesh_paths: Dict[str, str],
+        meshes: Dict[str, Mesh],
         batch_size: int = 1,
         device: torch.device = "cuda",
-        target_reduction: float = 0.0,
     ) -> None:
-        self._mesh_names = []
+        self._names = []
         self._vertices = []
         self._faces = []
         self._per_mesh_vertex_count = OrderedDict()
@@ -47,8 +46,8 @@ def __init__(
         self._lower_face_index_lookup = {}
         self._upper_face_index_lookup = {}
 
-        # load meshes
-        self._populate_meshes(mesh_paths, device, target_reduction)
+        # populate this container
+        self._populate_container(meshes, device)
 
         # add batch dim
         self._batch_size = batch_size
@@ -60,42 +59,22 @@ def __init__(
         self._device = device
 
     @abc.abstractmethod
-    def _load_mesh(self, mesh_path: str) -> trimesh.Trimesh:
-        return trimesh.load(mesh_path)
-
-    @abc.abstractmethod
-    def _populate_meshes(
+    def _populate_container(
         self,
-        mesh_paths: Dict[str, str],
+        meshes: Dict[str, Mesh],
         device: torch.device = "cuda",
-        target_reduction: float = 0.0,
     ) -> None:
         offset = 0
-        for mesh_name, mesh_path in mesh_paths.items():
+        for name, mesh in meshes.items():
             # populate mesh names
-            self._mesh_names.append(mesh_name)
-
-            # load mesh
-            m = self._load_mesh(mesh_path)
-
-            if isinstance(m, trimesh.Scene):
-                m = m.dump(concatenate=True)
-
-            vertices = m.vertices
-            faces = m.faces
-
-            vertices, faces = fast_simplification.simplify(
-                points=vertices,
-                triangles=faces,
-                target_reduction=target_reduction,
-            )
+            self._names.append(name)
 
             # populate mesh vertex count
-            self._per_mesh_vertex_count[mesh_name] = len(vertices)
+            self._per_mesh_vertex_count[name] = len(mesh.vertices)
 
             # populate vertices
             self._vertices.append(
-                torch.tensor(vertices, dtype=torch.float32, device=device)
+                torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
             )
             self._vertices[-1] = torch.cat(
                 [
@@ -106,22 +85,22 @@ def _populate_meshes(
             )  # (x,y,z) -> (x,y,z,1): homogeneous coordinates
 
             # populate mesh face count
-            self._per_mesh_face_count[mesh_name] = len(faces)
+            self._per_mesh_face_count[name] = len(mesh.faces)
 
             # populate faces (also add an offset to the point ids)
             self._faces.append(
                 torch.add(
-                    torch.tensor(faces, dtype=torch.int32, device=device),
+                    torch.tensor(mesh.faces, dtype=torch.int32, device=device),
                     offset,
                 )
             )
-            offset += len(vertices)
+            offset += len(mesh.vertices)
 
         self._vertices = torch.cat(self._vertices, dim=0)
         self._faces = torch.cat(self._faces, dim=0)
 
     def _populate_index_lookups(self) -> None:
-        if len(self._mesh_names) == 0:
+        if len(self._names) == 0:
             raise ValueError("No meshes loaded.")
         if len(self._per_mesh_vertex_count) == 0:
             raise ValueError("No vertex counts populated.")
@@ -131,16 +110,16 @@ def _populate_index_lookups(self) -> None:
         # crucial: self._per_mesh_vertex_count sorted same as self._vertices! Same for faces.
         running_vertex_index = 0
         running_face_index = 0
-        for mesh_name in self._mesh_names:
+        for name in self._names:
             # vertex index lookup
-            self._lower_vertex_index_lookup[mesh_name] = running_vertex_index
-            running_vertex_index += self._per_mesh_vertex_count[mesh_name]
-            self._upper_vertex_index_lookup[mesh_name] = running_vertex_index
+            self._lower_vertex_index_lookup[name] = running_vertex_index
+            running_vertex_index += self._per_mesh_vertex_count[name]
+            self._upper_vertex_index_lookup[name] = running_vertex_index
 
             # face index lookup
-            self._lower_face_index_lookup[mesh_name] = running_face_index
-            running_face_index += self._per_mesh_face_count[mesh_name]
-            self._upper_face_index_lookup[mesh_name] = running_face_index
+            self._lower_face_index_lookup[name] = running_face_index
+            running_face_index += self._per_mesh_face_count[name]
+            self._upper_face_index_lookup[name] = running_face_index
 
     @property
     def vertices(self) -> torch.FloatTensor:
@@ -171,8 +150,8 @@ def upper_face_index_lookup(self) -> Dict[str, int]:
         return self._upper_face_index_lookup
 
     @property
-    def mesh_names(self) -> List[str]:
-        return self._mesh_names
+    def names(self) -> List[str]:
+        return self._names
 
     @property
     def device(self) -> torch.device:

roboreg/io/__init__.py

@@ -1,3 +1,4 @@
 from .datasets import *
 from .filesystem import *
+from .meshes import *
 from .parsers import *

roboreg/io/meshes.py

@@ -0,0 +1,94 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Union
+
+import fast_simplification
+import numpy as np
+import trimesh
+
+
+@dataclass
+class Mesh:
+    r"""Dataclass to hold mesh data."""
+
+    vertices: np.ndarray
+    faces: np.ndarray
+
+
+def load_mesh(path: Union[Path, str]) -> Mesh:
+    r"""Reads a mesh file and returns vertices and faces.
+
+    Args:
+        path (Union[Path, str]): Path to the mesh file.
+
+    Returns:
+        Mesh:
+            - Vertices of shape Nx3.
+            - Faces of shape Nx3.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Mesh file {path} does not exist.")
+    m = trimesh.load(path)
+    if isinstance(m, trimesh.Scene):
+        m = m.dump(concatenate=True)
+    vertices, faces = m.vertices, m.faces
+    if vertices.size == 0 or faces.size == 0:
+        raise ValueError(f"Mesh is empty: {path.name}")
+    return Mesh(vertices=vertices.view(np.ndarray), faces=faces.view(np.ndarray))
+
+
+def load_meshes(
+    paths: Dict[str, Union[Path, str]],
+) -> Dict[str, Mesh]:
+    r"""Load multiple meshes.
+
+    Args:
+        paths (Dict[str, Union[Path, str]]): Mesh names and corresponding paths.
+
+    Returns:
+        Dict[str, Mesh]:
+            - Mesh name.
+            - Mesh vertices of shape Nx3 and faces of shape Nx3.
+    """
+    return {name: load_mesh(path) for name, path in paths.items()}
+
+
+def simplify_mesh(mesh: Mesh, target_reduction: float = 0.0) -> Mesh:
+    r"""Simplify a mesh.
+
+    Args:
+        mesh (Mesh): The mesh to be simplified.
+        target_reduction (float): Target reduction in [0, 1]. Zero for no reduction.
+
+    Returns:
+        Mesh: The simplified mesh.
+    """
+    if 0.0 > target_reduction > 1.0:
+        raise ValueError(
+            f"Expected target reduction in [0, 1], got {target_reduction}."
+        )
+    vertices, faces = fast_simplification.simplify(
+        points=mesh.vertices,
+        triangles=mesh.faces,
+        target_reduction=target_reduction,
+    )
+    return Mesh(vertices=vertices, faces=faces)
+
+
+def simplify_meshes(
+    meshes: Dict[str, Mesh], target_reduction: float = 0.0
+) -> Dict[str, Mesh]:
+    f"""Simplify multiple meshes.
+
+    Args:
+        meshes (Dict[str, Mesh]): The meshes to be simplified.
+        target_reduction (float): The target reduction in [0, 1]. Zero for no reduction.
+
+    Returns:
+        Dict[str, Mesh]: The simplified meshes. 
+    """
+    return {
+        name: simplify_mesh(mesh, target_reduction=target_reduction)
+        for name, mesh in meshes.items()
+    }

roboreg/util/sampling.py

@@ -1,6 +1,5 @@
 from typing import Tuple
 
-import numpy as np
 import torch
 
 

test/differentiable/test_kinematics.py

@@ -1,10 +1,3 @@
-import os
-import sys
-
-sys.path.append(
-    os.path.dirname((os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
-)
-
 import cv2
 import numpy as np
 import pytest
@@ -15,7 +8,7 @@
 from roboreg.differentiable.kinematics import TorchKinematics
 from roboreg.differentiable.rendering import NVDiffRastRenderer
 from roboreg.differentiable.structs import TorchMeshContainer
-from roboreg.io import URDFParser
+from roboreg.io import URDFParser, load_meshes
 from roboreg.util import from_homogeneous
 
 
@@ -60,7 +53,10 @@ def test_torch_kinematics_on_mesh(
         device=device,
     )
     meshes = TorchMeshContainer(
-        urdf_parser.ros_package_mesh_paths(root_link_name, end_link_name), device=device
+        meshes=load_meshes(
+            urdf_parser.ros_package_mesh_paths(root_link_name, end_link_name)
+        ),
+        device=device,
     )
 
     # compute forward kinematics and apply transforms to the meshes
@@ -112,7 +108,10 @@ def test_diff_kinematics() -> None:
         device=device,
     )
     meshes = TorchMeshContainer(
-        urdf_parser.ros_package_mesh_paths("lbr_link_0", "lbr_link_7"), device=device
+        meshes=load_meshes(
+            urdf_parser.ros_package_mesh_paths("lbr_link_0", "lbr_link_7")
+        ),
+        device=device,
     )
 
     # compute forward kinematics and apply transforms to the meshes
@@ -220,22 +219,28 @@ def test_diff_kinematics() -> None:
 
 
 if __name__ == "__main__":
-    # test_torch_kinematics(
-    #     ros_package="lbr_description",
-    #     xacro_path="urdf/med7/med7.xacro",
-    #     root_link_name="lbr_link_0",
-    #     end_link_name="lbr_link_7",
-    # )
-    test_torch_kinematics_on_mesh(
+    import os
+    import sys
+
+    os.environ["QT_QPA_PLATFORM"] = "offscreen"
+    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+    test_torch_kinematics(
         ros_package="lbr_description",
         xacro_path="urdf/med7/med7.xacro",
         root_link_name="lbr_link_0",
         end_link_name="lbr_link_7",
     )
     test_torch_kinematics_on_mesh(
-        ros_package="xarm_description",
-        xacro_path="urdf/xarm_device.urdf.xacro",
-        root_link_name="link_base",
-        end_link_name="link7",
+        ros_package="lbr_description",
+        xacro_path="urdf/med7/med7.xacro",
+        root_link_name="lbr_link_0",
+        end_link_name="lbr_link_7",
     )
+    # test_torch_kinematics_on_mesh(
+    #     ros_package="xarm_description",
+    #     xacro_path="urdf/xarm_device.urdf.xacro",
+    #     root_link_name="link_base",
+    #     end_link_name="link7",
+    # )
     # test_diff_kinematics()