optimize.py

import numpy as np
from scipy.spatial.transform import Rotation as R

import torch
from torch.optim.adam import Adam
import open3d

import lietorch


def transform_points(points: np.ndarray, transform: np.ndarray) -> np.ndarray:
    return np.dot(transform[:3, :3], points.T).T + transform[:3, 3]


def get_random_transformation():
    T = np.eye(4)
    T[:3, :3] = R.random().as_matrix()
    T[:3, 3] = 1e2 * (np.random.randn(3) - 1 / 2)
    return T


def optimize(
    p1, p2s, max_iter=int(5e3)
) -> tuple[torch.Tensor, list[np.ndarray], float]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    M = len(p2s)

    # X0 must be optimized
    X0 = torch.tensor(p1, dtype=torch.float32, requires_grad=True, device=device)
    Xs = [torch.tensor(p, dtype=torch.float32) for p in p2s]

    # represent the transformations as LieGroupParameters so they're on the manifold
    Ts = [
        lietorch.LieGroupParameter(lietorch.SE3.Identity(1, dtype=torch.float32))
        for _ in range(M)
    ]

    optimizer = Adam([X0, *Ts], lr=1e-1)

    ITER_WIDTH = len(str(max_iter))
    for i in range(max_iter):
        if i % 3000 == 0:
            for param_group in optimizer.param_groups:
                param_group["lr"] *= 0.75
        optimizer.zero_grad()

        err: torch.Tensor = torch.zeros(1).to(device)
        # todo: vectorize?
        for m in range(M):
            T = Ts[m]
            T_mat = T.retr().matrix()[0].to(device)

            # Transform X0 using the learned transformation
            X0_pred = torch.matmul(X0, T_mat[:3, :3].T) + T_mat[:3, 3]
            residual = X0_pred - Xs[m].to(device)
            err += torch.norm(residual, p=2).square()

        # backprop
        err.backward()
        optimizer.step()

        if i % 500 == 0 or (i + 1) == max_iter:
            print(f"iter {i:{ITER_WIDTH}} -> err: {err.item():.3f}")

    pred_Ts = [T.retr().matrix()[0].detach().numpy() for T in Ts]
    return X0, pred_Ts, err.item()


def icp_slam(M: int, position_noise_var: float = 1e1):
    assert M > 0, "number of additional frames should be more than 0"

    X0 = np.asarray(open3d.io.read_point_cloud("toothless.ply").points)
    T = [get_random_transformation() for _ in range(M)]
    X = [transform_points(X0, T[m]) for m in range(M)]

    X0_noisy = X0 + np.random.normal(0, position_noise_var, X0.shape)
    X_noisy = [x + np.random.normal(0, position_noise_var, x.shape) for x in X]

    _X0_pred, _Ts, err = optimize(X0_noisy, X_noisy)
    print("final error: ", err)


if __name__ == "__main__":
    icp_slam(5)