Add Volume-Based Centroid computation to structure/meshes

pytorch3d/ops/watertight.py

@@ -0,0 +1,33 @@
+import torch
+
+
+def volume_centroid(mesh):
+    """
+    Compute the volumetric centroid of this mesh, which is distinct from the center of mass.
+    The center of mass (average of all vertices) will be closer to where there are a
+    higher density of points in a mesh are, but the centroid, which is based on volume,
+    will be closer to a perceived center of the mesh, as opposed to based on the density
+    of vertices. This function assumes that the mesh is watertight, and that the faces are
+    all oriented in the same direction.
+    Returns:
+      The position of the centroid as a tensor of shape (3).
+    """
+    v_idxs = mesh.faces_padded().split([1, 1, 1], dim=-1)
+    verts = mesh.verts_padded()
+    valid = (mesh.faces_padded() != -1).all(dim=-1, keepdim=True)
+
+    v0, v1, v2 = [
+        torch.gather(
+            verts,
+            1,
+            idx.where(valid, torch.zeros_like(idx)).expand(-1, -1, 3),
+        ).where(valid, torch.zeros_like(idx, dtype=verts.dtype))
+        for idx in v_idxs
+    ]
+
+    tetra_center = (v0 + v1 + v2) / 4
+    signed_tetra_vol = (v0 * torch.cross(v1, v2, dim=-1)).sum(dim=-1, keepdim=True) / 6
+    denom = signed_tetra_vol.sum(dim=-2)
+    # clamp the denominator to prevent instability for degenerate meshes.
+    denom = torch.where(denom < 0, denom.clamp(max=-1e-5), denom.clamp(min=1e-5))
+    return (tetra_center * signed_tetra_vol).sum(dim=-2) / denom

tests/test_meshes.py

@@ -1298,6 +1298,25 @@ def test_assigned_normals(self):
             yes_normals.offset_verts_(torch.FloatTensor([1, 2, 3]).expand(12, 3))
             self.assertFalse(torch.allclose(yes_normals.verts_normals_padded(), verts))
 
+    def test_centroid(self):
+        meshes = init_simple_mesh()
+        # Check that it returns a valid value for multiple meshes with an inconsistent number
+        # of vertices
+        meshes.volume_centroid()
+
+        cube = init_cube_meshes()
+        self.assertClose(
+            cube.volume_centroid(),
+            torch.tensor(
+                [
+                    [0.5] * 3,
+                    [1.5] * 3,
+                    [2.5] * 3,
+                    [3.5] * 3,
+                ]
+            ),
+        )
+
     def test_submeshes(self):
         empty_mesh = Meshes([], [])
         # Four cubes with offsets [0, 1, 2, 3].