ICP - point-to-point version

Summary:
The iterative closest point algorithm - point-to-point version.

Output of `bm_iterative_closest_point`:
Argument key: `batch_size dim n_points_X n_points_Y use_pointclouds`

```
Benchmark                                         Avg Time(μs)      Peak Time(μs) Iterations
--------------------------------------------------------------------------------
IterativeClosestPoint_1_3_100_100_False              107569          111323              5
IterativeClosestPoint_1_3_100_1000_False             118972          122306              5
IterativeClosestPoint_1_3_1000_100_False             108576          110978              5
IterativeClosestPoint_1_3_1000_1000_False            331836          333515              2
IterativeClosestPoint_1_20_100_100_False             134387          137842              4
IterativeClosestPoint_1_20_100_1000_False            149218          153405              4
IterativeClosestPoint_1_20_1000_100_False            414248          416595              2
IterativeClosestPoint_1_20_1000_1000_False           374318          374662              2
IterativeClosestPoint_10_3_100_100_False             539852          539852              1
IterativeClosestPoint_10_3_100_1000_False            752784          752784              1
IterativeClosestPoint_10_3_1000_100_False           1070700         1070700              1
IterativeClosestPoint_10_3_1000_1000_False          1164020         1164020              1
IterativeClosestPoint_10_20_100_100_False            374548          377337              2
IterativeClosestPoint_10_20_100_1000_False           472764          476685              2
IterativeClosestPoint_10_20_1000_100_False          1457175         1457175              1
IterativeClosestPoint_10_20_1000_1000_False         2195820         2195820              1
IterativeClosestPoint_1_3_100_100_True               110084          115824              5
IterativeClosestPoint_1_3_100_1000_True              142728          147696              4
IterativeClosestPoint_1_3_1000_100_True              212966          213966              3
IterativeClosestPoint_1_3_1000_1000_True             369130          375114              2
IterativeClosestPoint_10_3_100_100_True              354615          355179              2
IterativeClosestPoint_10_3_100_1000_True             451815          452704              2
IterativeClosestPoint_10_3_1000_100_True             511833          511833              1
IterativeClosestPoint_10_3_1000_1000_True            798453          798453              1
--------------------------------------------------------------------------------
```

Reviewed By: shapovalov, gkioxari

Differential Revision: D19909952

fbshipit-source-id: f77fadc88fb7c53999909d594114b182ee2a3def

Author: davnov134

pytorch3d/ops/__init__.py

@@ -6,9 +6,10 @@
 from .knn import knn_gather, knn_points
 from .mesh_face_areas_normals import mesh_face_areas_normals
 from .packed_to_padded import packed_to_padded, padded_to_packed
-from .points_alignment import corresponding_points_alignment
+from .points_alignment import corresponding_points_alignment, iterative_closest_point
 from .sample_points_from_meshes import sample_points_from_meshes
 from .subdivide_meshes import SubdivideMeshes
+from .utils import convert_pointclouds_to_tensor, eyes, is_pointclouds, wmean
 from .vert_align import vert_align
 
 

pytorch3d/ops/points_alignment.py

@@ -1,22 +1,231 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
 import warnings
-from typing import List, Tuple, Union
+from typing import TYPE_CHECKING, List, NamedTuple, Optional, Union
 
 import torch
-from pytorch3d.ops import utils as oputil
+from pytorch3d.ops import knn_points
 from pytorch3d.structures import utils as strutil
-from pytorch3d.structures.pointclouds import Pointclouds
+
+from . import utils as oputil
+
+
+if TYPE_CHECKING:
+    from pytorch3d.structures.pointclouds import Pointclouds
+
+
+# named tuples for inputs/outputs
+class SimilarityTransform(NamedTuple):
+    R: torch.Tensor
+    T: torch.Tensor
+    s: torch.Tensor
+
+
+class ICPSolution(NamedTuple):
+    converged: bool
+    rmse: Union[torch.Tensor, None]
+    Xt: torch.Tensor
+    RTs: SimilarityTransform
+    t_history: List[SimilarityTransform]
+
+
+def iterative_closest_point(
+    X: Union[torch.Tensor, "Pointclouds"],
+    Y: Union[torch.Tensor, "Pointclouds"],
+    init_transform: Optional[SimilarityTransform] = None,
+    max_iterations: int = 100,
+    relative_rmse_thr: float = 1e-6,
+    estimate_scale: bool = False,
+    allow_reflection: bool = False,
+    verbose: bool = False,
+) -> ICPSolution:
+    """
+    Executes the iterative closest point (ICP) algorithm [1, 2] in order to find
+    a similarity transformation (rotation `R`, translation `T`, and
+    optionally scale `s`) between two given differently-sized sets of
+    `d`-dimensional points `X` and `Y`, such that:
+
+    `s[i] X[i] R[i] + T[i] = Y[NN[i]]`,
+
+    for all batch indices `i` in the least squares sense. Here, Y[NN[i]] stands
+    for the indices of nearest neighbors from `Y` to each point in `X`.
+    Note, however, that the solution is only a local optimum.
+
+    Args:
+        **X**: Batch of `d`-dimensional points
+            of shape `(minibatch, num_points_X, d)` or a `Pointclouds` object.
+        **Y**: Batch of `d`-dimensional points
+            of shape `(minibatch, num_points_Y, d)` or a `Pointclouds` object.
+        **init_transform**: A named-tuple `SimilarityTransform` of tensors
+            `R`, `T, `s`, where `R` is a batch of orthonormal matrices of
+            shape `(minibatch, d, d)`, `T` is a batch of translations
+            of shape `(minibatch, d)` and `s` is a batch of scaling factors
+            of shape `(minibatch,)`.
+        **max_iterations**: The maximum number of ICP iterations.
+        **relative_rmse_thr**: A threshold on the relative root mean squared error
+            used to terminate the algorithm.
+        **estimate_scale**: If `True`, also estimates a scaling component `s`
+            of the transformation. Otherwise assumes the identity
+            scale and returns a tensor of ones.
+        **allow_reflection**: If `True`, allows the algorithm to return `R`
+            which is orthonormal but has determinant==-1.
+        **verbose**: If `True`, prints status messages during each ICP iteration.
+
+    Returns:
+        A named tuple `ICPSolution` with the following fields:
+        **converged**: A boolean flag denoting whether the algorithm converged
+            successfully (=`True`) or not (=`False`).
+        **rmse**: Attained root mean squared error after termination of ICP.
+        **Xt**: The point cloud `X` transformed with the final transformation
+            (`R`, `T`, `s`). If `X` is a `Pointclouds` object, returns an
+            instance of `Pointclouds`, otherwise returns `torch.Tensor`.
+        **RTs**: A named tuple `SimilarityTransform` containing
+        a batch of similarity transforms with fields:
+            **R**: Batch of orthonormal matrices of shape `(minibatch, d, d)`.
+            **T**: Batch of translations of shape `(minibatch, d)`.
+            **s**: batch of scaling factors of shape `(minibatch, )`.
+        **t_history**: A list of named tuples `SimilarityTransform`
+            the transformation parameters after each ICP iteration.
+
+    References:
+        [1] Besl & McKay: A Method for Registration of 3-D Shapes. TPAMI, 1992.
+        [2] https://en.wikipedia.org/wiki/Iterative_closest_point
+    """
+
+    # make sure we convert input Pointclouds structures to
+    # padded tensors of shape (N, P, 3)
+    Xt, num_points_X = oputil.convert_pointclouds_to_tensor(X)
+    Yt, num_points_Y = oputil.convert_pointclouds_to_tensor(Y)
+
+    b, size_X, dim = Xt.shape
+
+    if (Xt.shape[2] != Yt.shape[2]) or (Xt.shape[0] != Yt.shape[0]):
+        raise ValueError(
+            "Point sets X and Y have to have the same "
+            + "number of batches and data dimensions."
+        )
+
+    if ((num_points_Y < Yt.shape[1]).any() or (num_points_X < Xt.shape[1]).any()) and (
+        num_points_Y != num_points_X
+    ).any():
+        # we have a heterogeneous input (e.g. because X/Y is
+        # an instance of Pointclouds)
+        mask_X = (
+            torch.arange(size_X, dtype=torch.int64, device=Xt.device)[None]
+            < num_points_X[:, None]
+        ).type_as(Xt)
+    else:
+        mask_X = Xt.new_ones(b, size_X)
+
+    # clone the initial point cloud
+    Xt_init = Xt.clone()
+
+    if init_transform is not None:
+        # parse the initial transform from the input and apply to Xt
+        try:
+            R, T, s = init_transform
+            assert (
+                R.shape == torch.Size((b, dim, dim))
+                and T.shape == torch.Size((b, dim))
+                and s.shape == torch.Size((b,))
+            )
+        except Exception:
+            raise ValueError(
+                "The initial transformation init_transform has to be "
+                "a named tuple SimilarityTransform with elements (R, T, s). "
+                "R are dim x dim orthonormal matrices of shape "
+                "(minibatch, dim, dim), T is a batch of dim-dimensional "
+                "translations of shape (minibatch, dim) and s is a batch "
+                "of scalars of shape (minibatch,)."
+            )
+        # apply the init transform to the input point cloud
+        Xt = _apply_similarity_transform(Xt, R, T, s)
+    else:
+        # initialize the transformation with identity
+        R = oputil.eyes(dim, b, device=Xt.device, dtype=Xt.dtype)
+        T = Xt.new_zeros((b, dim))
+        s = Xt.new_ones(b)
+
+    prev_rmse = None
+    rmse = None
+    iteration = -1
+    converged = False
+
+    # initialize the transformation history
+    t_history = []
+
+    # the main loop over ICP iterations
+    for iteration in range(max_iterations):
+        Xt_nn_points = knn_points(
+            Xt, Yt, lengths1=num_points_X, lengths2=num_points_Y, K=1, return_nn=True
+        )[2][:, :, 0, :]
+
+        # get the alignment of the nearest neighbors from Yt with Xt_init
+        R, T, s = corresponding_points_alignment(
+            Xt_init,
+            Xt_nn_points,
+            weights=mask_X,
+            estimate_scale=estimate_scale,
+            allow_reflection=allow_reflection,
+        )
+
+        # apply the estimated similarity transform to Xt_init
+        Xt = _apply_similarity_transform(Xt_init, R, T, s)
+
+        # add the current transformation to the history
+        t_history.append(SimilarityTransform(R, T, s))
+
+        # compute the root mean squared error
+        Xt_sq_diff = ((Xt - Xt_nn_points) ** 2).sum(2)
+        rmse = oputil.wmean(Xt_sq_diff[:, :, None], mask_X).sqrt()[:, 0, 0]
+
+        # compute the relative rmse
+        if prev_rmse is None:
+            relative_rmse = rmse.new_ones(b)
+        else:
+            relative_rmse = (prev_rmse - rmse) / prev_rmse
+
+        if verbose:
+            rmse_msg = (
+                f"ICP iteration {iteration}: mean/max rmse = "
+                + f"{rmse.mean():1.2e}/{rmse.max():1.2e} "
+                + f"; mean relative rmse = {relative_rmse.mean():1.2e}"
+            )
+            print(rmse_msg)
+
+        # check for convergence
+        if (relative_rmse <= relative_rmse_thr).all():
+            converged = True
+            break
+
+        # update the previous rmse
+        prev_rmse = rmse
+
+    if verbose:
+        if converged:
+            print(f"ICP has converged in {iteration + 1} iterations.")
+        else:
+            print(f"ICP has not converged in {max_iterations} iterations.")
+
+    if oputil.is_pointclouds(X):
+        Xt = X.update_padded(Xt)  # type: ignore
+
+    return ICPSolution(converged, rmse, Xt, SimilarityTransform(R, T, s), t_history)
+
+
+# threshold for checking that point crosscorelation
+# is full rank in corresponding_points_alignment
+AMBIGUOUS_ROT_SINGULAR_THR = 1e-15
 
 
 def corresponding_points_alignment(
-    X: Union[torch.Tensor, Pointclouds],
-    Y: Union[torch.Tensor, Pointclouds],
+    X: Union[torch.Tensor, "Pointclouds"],
+    Y: Union[torch.Tensor, "Pointclouds"],
     weights: Union[torch.Tensor, List[torch.Tensor], None] = None,
     estimate_scale: bool = False,
     allow_reflection: bool = False,
-    eps: float = 1e-8,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    eps: float = 1e-9,
+) -> SimilarityTransform:
     """
     Finds a similarity transformation (rotation `R`, translation `T`
     and optionally scale `s`)  between two given sets of corresponding
@@ -29,25 +238,25 @@ def corresponding_points_alignment(
     The algorithm is also known as Umeyama [1].
 
     Args:
-        X: Batch of `d`-dimensional points of shape `(minibatch, num_point, d)`
+        **X**: Batch of `d`-dimensional points of shape `(minibatch, num_point, d)`
             or a `Pointclouds` object.
-        Y: Batch of `d`-dimensional points of shape `(minibatch, num_point, d)`
+        **Y**: Batch of `d`-dimensional points of shape `(minibatch, num_point, d)`
             or a `Pointclouds` object.
-        weights: Batch of non-negative weights of
+        **weights**: Batch of non-negative weights of
             shape `(minibatch, num_point)` or list of `minibatch` 1-dimensional
             tensors that may have different shapes; in that case, the length of
             i-th tensor should be equal to the number of points in X_i and Y_i.
             Passing `None` means uniform weights.
-        estimate_scale: If `True`, also estimates a scaling component `s`
+        **estimate_scale**: If `True`, also estimates a scaling component `s`
             of the transformation. Otherwise assumes an identity
             scale and returns a tensor of ones.
-        allow_reflection: If `True`, allows the algorithm to return `R`
+        **allow_reflection**: If `True`, allows the algorithm to return `R`
             which is orthonormal but has determinant==-1.
-        eps: A scalar for clamping to avoid dividing by zero. Active for the
+        **eps**: A scalar for clamping to avoid dividing by zero. Active for the
             code that estimates the output scale `s`.
 
     Returns:
-        3-element tuple containing
+        3-element named tuple `SimilarityTransform` containing
         - **R**: Batch of orthonormal matrices of shape `(minibatch, d, d)`.
         - **T**: Batch of translations of shape `(minibatch, d)`.
         - **s**: batch of scaling factors of shape `(minibatch, )`.
@@ -58,8 +267,8 @@ def corresponding_points_alignment(
     """
 
     # make sure we convert input Pointclouds structures to tensors
-    Xt, num_points = _convert_point_cloud_to_tensor(X)
-    Yt, num_points_Y = _convert_point_cloud_to_tensor(Y)
+    Xt, num_points = oputil.convert_pointclouds_to_tensor(X)
+    Yt, num_points_Y = oputil.convert_pointclouds_to_tensor(Y)
 
     if (Xt.shape != Yt.shape) or (num_points != num_points_Y).any():
         raise ValueError(
@@ -90,8 +299,8 @@ def corresponding_points_alignment(
         weights = mask if weights is None else mask * weights.type_as(Xt)
 
     # compute the centroids of the point sets
-    Xmu = oputil.wmean(Xt, weights, eps=eps)
-    Ymu = oputil.wmean(Yt, weights, eps=eps)
+    Xmu = oputil.wmean(Xt, weight=weights, eps=eps)
+    Ymu = oputil.wmean(Yt, weight=weights, eps=eps)
 
     # mean-center the point sets
     Xc = Xt - Xmu
@@ -107,7 +316,7 @@ def corresponding_points_alignment(
     if (num_points < (dim + 1)).any():
         warnings.warn(
             "The size of one of the point clouds is <= dim+1. "
-            + "corresponding_points_alignment can't return a unique solution."
+            + "corresponding_points_alignment cannot return a unique rotation."
         )
 
     # compute the covariance XYcov between the point sets Xc, Yc
@@ -117,6 +326,16 @@ def corresponding_points_alignment(
     # decompose the covariance matrix XYcov
     U, S, V = torch.svd(XYcov)
 
+    # catch ambiguous rotation by checking the magnitude of singular values
+    if (S.abs() <= AMBIGUOUS_ROT_SINGULAR_THR).any() and not (
+        num_points < (dim + 1)
+    ).any():
+        warnings.warn(
+            "Excessively low rank of "
+            + "cross-correlation between aligned point clouds. "
+            + "corresponding_points_alignment cannot return a unique rotation."
+        )
+
     # identity matrix used for fixing reflections
     E = torch.eye(dim, dtype=XYcov.dtype, device=XYcov.device)[None].repeat(b, 1, 1)
 
@@ -148,26 +367,18 @@ def corresponding_points_alignment(
         # unit scaling since we do not estimate scale
         s = T.new_ones(b)
 
-    return R, T, s
+    return SimilarityTransform(R, T, s)
 
 
-def _convert_point_cloud_to_tensor(pcl: Union[torch.Tensor, Pointclouds]):
+def _apply_similarity_transform(
+    X: torch.Tensor, R: torch.Tensor, T: torch.Tensor, s: torch.Tensor
+) -> torch.Tensor:
     """
-    If `type(pcl)==Pointclouds`, converts a `pcl` object to a
-    padded representation and returns it together with the number of points
-    per batch. Otherwise, returns the input itself with the number of points
-    set to the size of the second dimension of `pcl`.
+    Applies a similarity transformation parametrized with a batch of orthonormal
+    matrices `R` of shape `(minibatch, d, d)`, a batch of translations `T`
+    of shape `(minibatch, d)` and a batch of scaling factors `s`
+    of shape `(minibatch,)` to a given `d`-dimensional cloud `X`
+    of shape `(minibatch, num_points, d)`
     """
-    if isinstance(pcl, Pointclouds):
-        X = pcl.points_padded()
-        num_points = pcl.num_points_per_cloud()
-    elif torch.is_tensor(pcl):
-        X = pcl
-        num_points = X.shape[1] * torch.ones(
-            X.shape[0], device=X.device, dtype=torch.int64
-        )
-    else:
-        raise ValueError(
-            "The inputs X, Y should be either Pointclouds objects or tensors."
-        )
-    return X, num_points
+    X = s[:, None, None] * torch.bmm(X, R) + T[:, None, :]
+    return X