diff --git a/pytorch3d/ops/watertight.py b/pytorch3d/ops/watertight.py new file mode 100644 index 000000000..ed5d73e16 --- /dev/null +++ b/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 diff --git a/pytorch3d/structures/meshes.py b/pytorch3d/structures/meshes.py index 97e0055a8..19f112641 100644 --- a/pytorch3d/structures/meshes.py +++ b/pytorch3d/structures/meshes.py @@ -1545,36 +1545,6 @@ def sample_textures(self, fragments): else: raise ValueError("Meshes does not have textures") - 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. 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_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 - 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, face_indices: Union[ diff --git a/tests/test_meshes.py b/tests/test_meshes.py index 74afe33f2..6f9f8de1f 100644 --- a/tests/test_meshes.py +++ b/tests/test_meshes.py @@ -1299,13 +1299,24 @@ def test_assigned_normals(self): 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, - ])) + 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].