|
| 1 | +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. |
| 2 | + |
| 3 | +from typing import TYPE_CHECKING, Tuple, Union |
| 4 | + |
| 5 | +import torch |
| 6 | + |
| 7 | +from .utils import convert_pointclouds_to_tensor, get_point_covariances |
| 8 | + |
| 9 | + |
| 10 | +if TYPE_CHECKING: |
| 11 | + from ..structures import Pointclouds |
| 12 | + |
| 13 | + |
| 14 | +def estimate_pointcloud_normals( |
| 15 | + pointclouds: Union[torch.Tensor, "Pointclouds"], |
| 16 | + neighborhood_size: int = 50, |
| 17 | + disambiguate_directions: bool = True, |
| 18 | +) -> torch.Tensor: |
| 19 | + """ |
| 20 | + Estimates the normals of a batch of `pointclouds`. |
| 21 | +
|
| 22 | + The function uses `estimate_pointcloud_local_coord_frames` to estimate |
| 23 | + the normals. Please refer to this function for more detailed information. |
| 24 | +
|
| 25 | + Args: |
| 26 | + **pointclouds**: Batch of 3-dimensional points of shape |
| 27 | + `(minibatch, num_point, 3)` or a `Pointclouds` object. |
| 28 | + **neighborhood_size**: The size of the neighborhood used to estimate the |
| 29 | + geometry around each point. |
| 30 | + **disambiguate_directions**: If `True`, uses the algorithm from [1] to |
| 31 | + ensure sign consistency of the normals of neigboring points. |
| 32 | +
|
| 33 | + Returns: |
| 34 | + **normals**: A tensor of normals for each input point |
| 35 | + of shape `(minibatch, num_point, 3)`. |
| 36 | + If `pointclouds` are of `Pointclouds` class, returns a padded tensor. |
| 37 | +
|
| 38 | + References: |
| 39 | + [1] Tombari, Salti, Di Stefano: Unique Signatures of Histograms for |
| 40 | + Local Surface Description, ECCV 2010. |
| 41 | + """ |
| 42 | + |
| 43 | + curvatures, local_coord_frames = estimate_pointcloud_local_coord_frames( |
| 44 | + pointclouds, |
| 45 | + neighborhood_size=neighborhood_size, |
| 46 | + disambiguate_directions=disambiguate_directions, |
| 47 | + ) |
| 48 | + |
| 49 | + # the normals correspond to the first vector of each local coord frame |
| 50 | + normals = local_coord_frames[:, :, :, 0] |
| 51 | + |
| 52 | + return normals |
| 53 | + |
| 54 | + |
| 55 | +def estimate_pointcloud_local_coord_frames( |
| 56 | + pointclouds: Union[torch.Tensor, "Pointclouds"], |
| 57 | + neighborhood_size: int = 50, |
| 58 | + disambiguate_directions: bool = True, |
| 59 | +) -> Tuple[torch.Tensor, torch.Tensor]: |
| 60 | + """ |
| 61 | + Estimates the principal directions of curvature (which includes normals) |
| 62 | + of a batch of `pointclouds`. |
| 63 | +
|
| 64 | + The algorithm first finds `neighborhood_size` nearest neighbors for each |
| 65 | + point of the point clouds, followed by obtaining principal vectors of |
| 66 | + covariance matrices of each of the point neighborhoods. |
| 67 | + The main principal vector corresponds to the normals, while the |
| 68 | + other 2 are the direction of the highest curvature and the 2nd highest |
| 69 | + curvature. |
| 70 | +
|
| 71 | + Note that each principal direction is given up to a sign. Hence, |
| 72 | + the function implements `disambiguate_directions` switch that allows |
| 73 | + to ensure consistency of the sign of neighboring normals. The implementation |
| 74 | + follows the sign disabiguation from SHOT descriptors [1]. |
| 75 | +
|
| 76 | + The algorithm also returns the curvature values themselves. |
| 77 | + These are the eigenvalues of the estimated covariance matrices |
| 78 | + of each point neighborhood. |
| 79 | +
|
| 80 | + Args: |
| 81 | + **pointclouds**: Batch of 3-dimensional points of shape |
| 82 | + `(minibatch, num_point, 3)` or a `Pointclouds` object. |
| 83 | + **neighborhood_size**: The size of the neighborhood used to estimate the |
| 84 | + geometry around each point. |
| 85 | + **disambiguate_directions**: If `True`, uses the algorithm from [1] to |
| 86 | + ensure sign consistency of the normals of neigboring points. |
| 87 | +
|
| 88 | + Returns: |
| 89 | + **curvatures**: The three principal curvatures of each point |
| 90 | + of shape `(minibatch, num_point, 3)`. |
| 91 | + If `pointclouds` are of `Pointclouds` class, returns a padded tensor. |
| 92 | + **local_coord_frames**: The three principal directions of the curvature |
| 93 | + around each point of shape `(minibatch, num_point, 3, 3)`. |
| 94 | + The principal directions are stored in columns of the output. |
| 95 | + E.g. `local_coord_frames[i, j, :, 0]` is the normal of |
| 96 | + `j`-th point in the `i`-th pointcloud. |
| 97 | + If `pointclouds` are of `Pointclouds` class, returns a padded tensor. |
| 98 | +
|
| 99 | + References: |
| 100 | + [1] Tombari, Salti, Di Stefano: Unique Signatures of Histograms for |
| 101 | + Local Surface Description, ECCV 2010. |
| 102 | + """ |
| 103 | + |
| 104 | + points_padded, num_points = convert_pointclouds_to_tensor(pointclouds) |
| 105 | + |
| 106 | + ba, N, dim = points_padded.shape |
| 107 | + if dim != 3: |
| 108 | + raise ValueError( |
| 109 | + "The pointclouds argument has to be of shape (minibatch, N, 3)" |
| 110 | + ) |
| 111 | + |
| 112 | + if (num_points <= neighborhood_size).any(): |
| 113 | + raise ValueError( |
| 114 | + "The neighborhood_size argument has to be" |
| 115 | + + " >= size of each of the point clouds." |
| 116 | + ) |
| 117 | + |
| 118 | + # undo global mean for stability |
| 119 | + # TODO: replace with tutil.wmean once landed |
| 120 | + pcl_mean = points_padded.sum(1) / num_points[:, None] |
| 121 | + points_centered = points_padded - pcl_mean[:, None, :] |
| 122 | + |
| 123 | + # get the per-point covariance and nearest neighbors used to compute it |
| 124 | + cov, knns = get_point_covariances(points_centered, num_points, neighborhood_size) |
| 125 | + |
| 126 | + # get the local coord frames as principal directions of |
| 127 | + # the per-point covariance |
| 128 | + # this is done with torch.symeig, which returns the |
| 129 | + # eigenvectors (=principal directions) in an ascending order of their |
| 130 | + # corresponding eigenvalues, while the smallest eigenvalue's eigenvector |
| 131 | + # corresponds to the normal direction |
| 132 | + curvatures, local_coord_frames = torch.symeig(cov, eigenvectors=True) |
| 133 | + |
| 134 | + # disambiguate the directions of individual principal vectors |
| 135 | + if disambiguate_directions: |
| 136 | + # disambiguate normal |
| 137 | + n = _disambiguate_vector_directions( |
| 138 | + points_centered, knns, local_coord_frames[:, :, :, 0] |
| 139 | + ) |
| 140 | + # disambiguate the main curvature |
| 141 | + z = _disambiguate_vector_directions( |
| 142 | + points_centered, knns, local_coord_frames[:, :, :, 2] |
| 143 | + ) |
| 144 | + # the secondary curvature is just a cross between n and z |
| 145 | + y = torch.cross(n, z, dim=2) |
| 146 | + # cat to form the set of principal directions |
| 147 | + local_coord_frames = torch.stack((n, y, z), dim=3) |
| 148 | + |
| 149 | + return curvatures, local_coord_frames |
| 150 | + |
| 151 | + |
| 152 | +def _disambiguate_vector_directions(pcl, knns, vecs): |
| 153 | + """ |
| 154 | + Disambiguates normal directions according to [1]. |
| 155 | +
|
| 156 | + References: |
| 157 | + [1] Tombari, Salti, Di Stefano: Unique Signatures of Histograms for |
| 158 | + Local Surface Description, ECCV 2010. |
| 159 | + """ |
| 160 | + # parse out K from the shape of knns |
| 161 | + K = knns.shape[2] |
| 162 | + # the difference between the mean of each neighborhood and |
| 163 | + # each element of the neighborhood |
| 164 | + df = knns - pcl[:, :, None] |
| 165 | + # projection of the difference on the principal direction |
| 166 | + proj = (vecs[:, :, None] * df).sum(3) |
| 167 | + # check how many projections are positive |
| 168 | + n_pos = (proj > 0).type_as(knns).sum(2, keepdim=True) |
| 169 | + # flip the principal directions where number of positive correlations |
| 170 | + flip = (n_pos < (0.5 * K)).type_as(knns) |
| 171 | + vecs = (1.0 - 2.0 * flip) * vecs |
| 172 | + return vecs |
0 commit comments