Add test and switch to padded instead of packed

pytorch3d/structures/meshes.py

@@ -1545,27 +1545,35 @@ def sample_textures(self, fragments):
         else:
             raise ValueError("Meshes does not have textures")
 
-    def centroid(self):
+    def volume_centroid(self):
         """
         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.
+        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 = self.faces_packed().split([1, 1, 1], dim=-1)
-        verts = self.verts_packed()
-        v0, v1, v2 = [verts[idx].squeeze(-2) for idx in v_idxs]
+        v_idxs = self.faces_padded().split([1, 1, 1], dim=-1)
+        verts = self.verts_padded()
+
+        v0, v1, v2 = [torch.gather(verts, 1, idx.expand(-1, -1, 3)) 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
-        return (tetra_center * signed_tetra_vol).sum(dim=-2) / signed_tetra_vol.sum(
-            dim=-2
-        ).clamp(min=1e-5)
+        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
 
     def submeshes(
         self,

tests/test_meshes.py

@@ -1298,6 +1298,14 @@ 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):
+        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].