61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
|
|
def read_pcl_tensor(path):
|
|
pcl_np = read_pcl(path)
|
|
return torch.from_numpy(pcl_np)
|
|
|
|
|
|
def align_pcl(Y, X, weight=None, fixed_scale=False):
|
|
"""align similarity transform to align X with Y using umeyama method
|
|
X' = s * R * X + t is aligned with Y
|
|
:param Y (*, N, 3) first trajectory
|
|
:param X (*, N, 3) second trajectory
|
|
:param weight (*, N, 1) optional weight of valid correspondences
|
|
:returns s (*, 1), R (*, 3, 3), t (*, 3)
|
|
"""
|
|
*dims, N, _ = Y.shape
|
|
N = torch.ones(*dims, 1, 1) * N
|
|
|
|
if weight is not None:
|
|
Y = Y * weight
|
|
X = X * weight
|
|
N = weight.sum(dim=-2, keepdim=True) # (*, 1, 1)
|
|
|
|
# subtract mean
|
|
my = Y.sum(dim=-2) / N[..., 0] # (*, 3)
|
|
mx = X.sum(dim=-2) / N[..., 0]
|
|
y0 = Y - my[..., None, :] # (*, N, 3)
|
|
x0 = X - mx[..., None, :]
|
|
|
|
if weight is not None:
|
|
y0 = y0 * weight
|
|
x0 = x0 * weight
|
|
|
|
# correlation
|
|
C = torch.matmul(y0.transpose(-1, -2), x0) / N # (*, 3, 3)
|
|
U, D, Vh = torch.linalg.svd(C) # (*, 3, 3), (*, 3), (*, 3, 3)
|
|
|
|
S = torch.eye(3).reshape(*(1,) * (len(dims)), 3, 3).repeat(*dims, 1, 1)
|
|
neg = torch.det(U) * torch.det(Vh.transpose(-1, -2)) < 0
|
|
S[neg, 2, 2] = -1
|
|
|
|
R = torch.matmul(U, torch.matmul(S, Vh)) # (*, 3, 3)
|
|
|
|
D = torch.diag_embed(D) # (*, 3, 3)
|
|
if fixed_scale:
|
|
s = torch.ones(*dims, 1, device=Y.device, dtype=torch.float32)
|
|
else:
|
|
var = torch.sum(torch.square(x0), dim=(-1, -2), keepdim=True) / N # (*, 1, 1)
|
|
s = (
|
|
torch.diagonal(torch.matmul(D, S), dim1=-2, dim2=-1).sum(
|
|
dim=-1, keepdim=True
|
|
)
|
|
/ var[..., 0]
|
|
) # (*, 1)
|
|
|
|
t = my - s * torch.matmul(R, mx[..., None])[..., 0] # (*, 3)
|
|
|
|
return s, R, t
|