LDDMM: how to register a mesh to a template?#

[1]:

import shutil
import string

import numpy as np
import pyvista as pv

import polpo.utils as putils
from polpo.plot.pyvista import RegisteredMeshesGifPlotter
from polpo.preprocessing.mesh.registration import RigidAlignment
from polpo.surface_mesh.deformetrica import (
    LddmmMetric,
    Point,
    RegistrationDir,
    ShootDir,
)
from polpo.surface_mesh.generation.blob import create_blob

[KeOps] Warning : CUDA was detected, but driver API could not be initialized. Switching to CPU only.

[2]:

RECOMPUTE = False  # in case registration has already been run

# NB: fix seed before setting it to False
np.random.seed(42)

STATIC_VIZ = True

if STATIC_VIZ:
    pv.set_jupyter_backend("static")

[3]:

OUTPUTS_DIR = putils.get_results_path() / "registration_blob_example"

if OUTPUTS_DIR.exists() and RECOMPUTE:
    shutil.rmtree(OUTPUTS_DIR)

OUTPUTS_DIR

[3]:

PosixPath('/home/luisfpereira/.polpo/results/registration_blob_example')

Generate meshes#

[4]:

bump_amp = 0.2

raw_meshes = [
    create_blob(resolution=10, bump_amp=bump_amp, n_bumps=5, smoothing_iter=10)
    for _ in range(2)
]

raw_meshes[0].points.shape

[4]:

(82, 3)

[5]:

meshes_ = raw_meshes
pl = pv.Plotter(border=False)

pl.add_mesh(
    meshes_[0],
    show_edges=True,
    color="red",
    opacity=0.5,
    label="source",
)
pl.add_mesh(
    meshes_[1],
    show_edges=True,
    color="green",
    opacity=0.5,
    label="target",
)
pl.add_legend()

pl.show()

../../../_images/_generated_notebooks_how_to_deformetrica_register_mesh_template_6_0.png

Preprocessing#

As we can see in the visualization, meshes are not rigid aligned. Preprocessing takes care of this kind of details, before applying LDDMM.

[6]:

prep_pipe = RigidAlignment(known_correspondences=True)

meshes = prep_pipe(raw_meshes)

[7]:

meshes_ = meshes
pl = pv.Plotter(border=False)

pl.add_mesh(
    meshes_[0],
    show_edges=True,
    color="red",
    opacity=0.5,
    label="source",
)
pl.add_mesh(
    meshes_[1],
    show_edges=True,
    color="green",
    opacity=0.5,
    label="target",
)

pl.add_legend()

pl.show()

../../../_images/_generated_notebooks_how_to_deformetrica_register_mesh_template_10_0.png

Registration#

[8]:

kernel_width = 2 * bump_amp
registration_kwargs = dict(
    kernel_width=kernel_width,
    regularisation=1.0,
    max_iter=2000,
    freeze_control_points=False,
    metric="varifold",
    tol=1e-16,
    attachment_kernel_width=bump_amp,
)

metric = LddmmMetric(OUTPUTS_DIR, **registration_kwargs)

[9]:

points = [
    Point(
        id_=string.ascii_uppercase[index],
        pv_surface=mesh,
        dirname=metric.dir_config.meshes_dir,
    )
    for index, mesh in enumerate(meshes)
]

base_point, point = points

[10]:

# NB: computational time changes very little for res \in [10, 20]
vec = metric.log(point, base_point)

Visualization#

Template and control points#

[11]:

pl = pv.Plotter()

pl.add_mesh(
    base_point.as_pv(),
    show_edges=True,
    label="source",
)

pl.add_points(
    vec.control_points().as_pv(),
    color="red",
    label="control points",
)

pl.add_legend()

pl.show()

../../../_images/_generated_notebooks_how_to_deformetrica_register_mesh_template_17_0.png

Target and reconstructed meshes#

[12]:

reg_dir = RegistrationDir.from_dirname(vec.dirname)

reconstructed = reg_dir.reconstructed()

[13]:

pl = pv.Plotter()

pl.add_mesh(
    reconstructed.as_pv(),
    show_edges=True,
    opacity=0.5,
    label="reconstructed",
)
pl.add_mesh(
    point.as_pv(),
    opacity=0.5,
    color="red",
    label="target",
)
pl.add_legend()

pl.show()

../../../_images/_generated_notebooks_how_to_deformetrica_register_mesh_template_20_0.png

Flow#

[14]:

flow_meshes = [pt.as_pv() for pt in reg_dir.flow()]

[15]:

pl = RegisteredMeshesGifPlotter()

pl.add_mesh(reconstructed.as_pv(), opacity=0.5)
pl.add_title("source->target")

pl.add_meshes(flow_meshes)

pl.close()

pl.show()

../../../_images/_generated_notebooks_how_to_deformetrica_register_mesh_template_23_0.png

[16]:

np.amax(np.abs(flow_meshes[-1].points - reconstructed.as_pv().points))

[16]:

np.float32(0.0)

Reconstruct flow by shooting#

[17]:

point_ = metric.exp(vec, base_point)

[18]:

shoot_dir = ShootDir.from_dirname(point_.vtk_path.parent)

flow_meshes_ = [pt.as_pv() for pt in shoot_dir.flow()]

[19]:

[
    np.amax(np.abs(flow_meshes_[0].points - base_point.as_pv().points)),
    np.amax(np.abs(flow_meshes_[-1].points - point_.as_pv().points)),
    np.amax(point.as_pv().points - point_.as_pv().points),
]

[19]:

[np.float64(2.8968215604052716e-08),
 np.float32(0.0),
 np.float64(0.1330844801730855)]

[20]:

diffs = [
    np.abs(np.amax(mesh.points - cmp_mesh.points))
    for mesh, cmp_mesh in zip(flow_meshes, flow_meshes_)
]

# TODO: identify origin of differences
np.amax(diffs)

[20]:

np.float32(0.014356971)

Further reading#