Compute a rigid transformation that, when applied to p, minimizes the weighted squared distance between transformed points in p and points in q. See "Least-Squares Rigid Motion Using SVD" by Olga Sorkine-Hornung and Michael Rabinovich for more details (https://igl.ethz.ch/projects/ARAP/s
(
p: Float[Tensor, "*batch point 3"],
q: Float[Tensor, "*batch point 3"],
weights: Float[Tensor, "*batch point"],
)
| 5 | |
| 6 | |
| 7 | def align_rigid( |
| 8 | p: Float[Tensor, "*batch point 3"], |
| 9 | q: Float[Tensor, "*batch point 3"], |
| 10 | weights: Float[Tensor, "*batch point"], |
| 11 | ) -> Float[Tensor, "*batch 4 4"]: |
| 12 | """Compute a rigid transformation that, when applied to p, minimizes the weighted |
| 13 | squared distance between transformed points in p and points in q. See "Least-Squares |
| 14 | Rigid Motion Using SVD" by Olga Sorkine-Hornung and Michael Rabinovich for more |
| 15 | details (https://igl.ethz.ch/projects/ARAP/svd_rot.pdf). |
| 16 | """ |
| 17 | |
| 18 | device = p.device |
| 19 | dtype = p.dtype |
| 20 | *batch, _, _ = p.shape |
| 21 | |
| 22 | # 1. Compute the centroids of both point sets. |
| 23 | weights_normalized = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8) |
| 24 | p_centroid = (weights_normalized[..., None] * p).sum(dim=-2) |
| 25 | q_centroid = (weights_normalized[..., None] * q).sum(dim=-2) |
| 26 | |
| 27 | # 2. Compute the centered vectors. |
| 28 | p_centered = p - p_centroid[..., None, :] |
| 29 | q_centered = q - q_centroid[..., None, :] |
| 30 | |
| 31 | # 3. Compute the 3x3 covariance matrix. |
| 32 | covariance = (q_centered * weights[..., None]).transpose(-1, -2) @ p_centered |
| 33 | |
| 34 | # 4. Compute the singular value decomposition and then the rotation. |
| 35 | u, _, vt = torch.linalg.svd(covariance) |
| 36 | s = torch.eye(3, dtype=dtype, device=device) |
| 37 | s = s.expand((*batch, 3, 3)).contiguous() |
| 38 | s[..., 2, 2] = (u.det() * vt.det()).sign() |
| 39 | rotation = u @ s @ vt |
| 40 | |
| 41 | # 5. Compute the optimal translation. |
| 42 | translation = q_centroid - einsum(rotation, p_centroid, "... i j, ... j -> ... i") |
| 43 | |
| 44 | # Compose the results into a single transformation matrix. |
| 45 | shape = (*rotation.shape[:-2], 4, 4) |
| 46 | r = torch.eye(4, dtype=torch.float32, device=device).expand(shape).contiguous() |
| 47 | r[..., :3, :3] = rotation |
| 48 | t = torch.eye(4, dtype=torch.float32, device=device).expand(shape).contiguous() |
| 49 | t[..., :3, 3] = translation |
| 50 | |
| 51 | return t @ r |