diff --git a/pytorch3d/ops/points_normals.py b/pytorch3d/ops/points_normals.py index e5b0ec337..702c0bb70 100644 --- a/pytorch3d/ops/points_normals.py +++ b/pytorch3d/ops/points_normals.py @@ -8,6 +8,7 @@ import torch +from ..common.workaround import symeig3x3 from .utils import convert_pointclouds_to_tensor, get_point_covariances @@ -19,6 +20,8 @@ def estimate_pointcloud_normals( pointclouds: Union[torch.Tensor, "Pointclouds"], neighborhood_size: int = 50, disambiguate_directions: bool = True, + *, + use_symeig_workaround: bool = True, ) -> torch.Tensor: """ Estimates the normals of a batch of `pointclouds`. @@ -33,6 +36,8 @@ def estimate_pointcloud_normals( geometry around each point. **disambiguate_directions**: If `True`, uses the algorithm from [1] to ensure sign consistency of the normals of neighboring points. + **use_symeig_workaround**: If `True`, uses a custom eigenvalue + calculation. Returns: **normals**: A tensor of normals for each input point @@ -48,6 +53,7 @@ def estimate_pointcloud_normals( pointclouds, neighborhood_size=neighborhood_size, disambiguate_directions=disambiguate_directions, + use_symeig_workaround=use_symeig_workaround, ) # the normals correspond to the first vector of each local coord frame @@ -60,6 +66,8 @@ def estimate_pointcloud_local_coord_frames( pointclouds: Union[torch.Tensor, "Pointclouds"], neighborhood_size: int = 50, disambiguate_directions: bool = True, + *, + use_symeig_workaround: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Estimates the principal directions of curvature (which includes normals) @@ -88,6 +96,8 @@ def estimate_pointcloud_local_coord_frames( geometry around each point. **disambiguate_directions**: If `True`, uses the algorithm from [1] to ensure sign consistency of the normals of neighboring points. + **use_symeig_workaround**: If `True`, uses a custom eigenvalue + calculation. Returns: **curvatures**: The three principal curvatures of each point @@ -133,7 +143,10 @@ def estimate_pointcloud_local_coord_frames( # eigenvectors (=principal directions) in an ascending order of their # corresponding eigenvalues, while the smallest eigenvalue's eigenvector # corresponds to the normal direction - curvatures, local_coord_frames = torch.symeig(cov, eigenvectors=True) + if use_symeig_workaround: + curvatures, local_coord_frames = symeig3x3(cov, eigenvectors=True) + else: + curvatures, local_coord_frames = torch.symeig(cov, eigenvectors=True) # disambiguate the directions of individual principal vectors if disambiguate_directions: diff --git a/tests/benchmarks/bm_points_normals.py b/tests/benchmarks/bm_points_normals.py new file mode 100644 index 000000000..487e71e54 --- /dev/null +++ b/tests/benchmarks/bm_points_normals.py @@ -0,0 +1,47 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import itertools + +import torch +from fvcore.common.benchmark import benchmark +from pytorch3d.ops import estimate_pointcloud_normals +from test_points_normals import TestPCLNormals + + +def to_bm(num_points, use_symeig_workaround): + device = torch.device("cuda:0") + points_padded, _normals = TestPCLNormals.init_spherical_pcl( + num_points=num_points, device=device, use_pointclouds=False + ) + torch.cuda.synchronize() + + def run(): + estimate_pointcloud_normals( + points_padded, use_symeig_workaround=use_symeig_workaround + ) + torch.cuda.synchronize() + + return run + + +def bm_points_normals() -> None: + case_grid = { + "use_symeig_workaround": [True, False], + "num_points": [3000, 6000], + } + test_cases = itertools.product(*case_grid.values()) + kwargs_list = [dict(zip(case_grid.keys(), case)) for case in test_cases] + benchmark( + to_bm, + "normals", + kwargs_list, + warmup_iters=1, + ) + + +if __name__ == "__main__": + bm_points_normals()