[Python] Shape registration using PCA and GPA not aligning

I am currently working on a statistical shape model/active shape model to make it fit a vertebra. But for some reason it is not working and I have been trying to fit it for some time.

First I tried fitting with landmarks. I took seven of the same vertebra and did annotations of 10 spots:



I aligned all meshes with the mean to 0,0,0 and rigidly registered all landmarks and meshes



Imports:

import os
import json
import glob
import trimesh
import numpy as np
from scipy.spatial import procrustes
from trimesh.registration import icp
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
from tqdm import tqdm


Then I did GPA and PCA:

def rigid_procrustes_with_scaling(X, Y):
"""
Procrustes alignment with scaling (similarity transform).
Returns Y_aligned, R, t, s
"""
X_mean = X.mean(axis=0)
Y_mean = Y.mean(axis=0)
X0 = X - X_mean
Y0 = Y - Y_mean

norm_X = np.linalg.norm(X0)
norm_Y = np.linalg.norm(Y0)
s = norm_X / norm_Y

U, _, Vt = np.linalg.svd((Y0 * s).T @ X0)
R = U @ Vt
if np.linalg.det(R) < 0:
Vt[-1, :] *= -1
R = U @ Vt

t = X_mean - (Y_mean * s) @ R
Y_aligned = (Y * s) @ R + t
return Y_aligned, R, t, s

def compute_landmark_gpa_pca(landmarks_dict, n_components=3, tol=1e-5, max_iter=100):
"""
landmarks_dict: dict of {name: (N,3) numpy array of landmarks}
"""
names = list(landmarks_dict.keys())
landmarks_list = [landmarks_dict[name] for name in names]
N = landmarks_list[0].shape[0]

for lm in landmarks_list:
if lm.shape != (N, 3):
raise ValueError("All landmark sets must have shape (N, 3)")

mean_shape = landmarks_list[0].copy()
for iteration in range(max_iter):
aligned = []
for lm in landmarks_list:
aligned_lm, _, _, _ = rigid_procrustes_with_scaling(mean_shape, lm)
aligned.append(aligned_lm)
new_mean = np.mean(aligned, axis=0)
diff = np.linalg.norm(mean_shape - new_mean)
if diff < tol:
print(f"GPA converged in {iteration+1} steps")
break
mean_shape = new_mean
else:
print("GPA did not fully converge.")

# PCA
flat_shapes = np.array([a.flatten() for a in aligned])
pca = PCA(n_components=n_components)
pca.fit(flat_shapes)

return mean_shape, pca


This is how I am fitting the PCA to a new mesh

def fit_active_shape_model_with_history(new_mesh, mean_shape, pca, n_iterations=10):
"""
Standard ASM:
- No scaling
- Rigid alignment
- PCA projection
Returns:
history: list of shape arrays per iteration (in model space)
b: final PCA weights
fitted_shape_aligned: final shape aligned to target mesh
"""
# Initialize
# Use similarity Procrustes for first alignment
closest_points, _, _ = trimesh.proximity.closest_point(new_mesh, mean_shape)
fitted_shape, R_init, t_init, s_init = rigid_procrustes_with_scaling(closest_points, mean_shape)

class PCASimple:
def __init__(self, mean_flat, components, explained_variance):
self.mean_ = mean_flat
self.components_ = components
self.explained_variance_ = explained_variance
def transform(self, X):
return (X - self.mean_) @ self.components_.T
def inverse_transform(self, b):
return self.mean_ + b @ self.components_

pca_model = PCASimple(mean_shape.flatten(), pca.components_, pca.explained_variance_)

history = []

for it in range(n_iterations):
# Find closest points
closest_points, _, _ = trimesh.proximity.closest_point(new_mesh, fitted_shape)
# Rigid alignment
fitted_shape, _, _, _ = rigid_procrustes_with_scaling(closest_points, fitted_shape)

# PCA projection
shape_flat = fitted_shape.flatten()
b = pca_model.transform(shape_flat)

stds = np.sqrt(pca_model.explained_variance_)
b = np.clip(b, -2 * stds, 2 * stds)

# Reconstruct
fitted_flat = pca_model.inverse_transform(b)
fitted_shape = fitted_flat.reshape(-1,3)

history.append(fitted_shape.copy())

# At the end: align fitted shape to target mesh
closest_points, _, _ = trimesh.proximity.closest_point(new_mesh, fitted_shape)
# Get aligned version + final transform
fitted_shape_aligned, R_final, t_final, s_final = rigid_procrustes_with_scaling(fitted_shape, closest_points)

return history, b, fitted_shape_aligned


But in total, everything is just off and doesn't make any sense.



I also tried the same but with surface point clouds to better recreate the shape of the vertebra (here I used a total of 40 vertebrae of the same type). The preprocessing is almost the same but I tried to get according points an all meshes (tried to do this by laying them over each other and taking the closest distance).

An image of all points after the preprocessing:



PCA and GPA is the same as above. This is my code for the fitting:

# ---------- TEMPLATE SAMPLING ----------
def create_template_points(template_mesh, n_points=1000, seed=42):
np.random.seed(seed)
return template_mesh.sample(n_points)

def sample_mesh_to_points(mesh, template_points):
closest_points, _, _ = trimesh.proximity.closest_point(mesh, template_points)
return closest_points

# ---------- PCA HELPER ----------
class PCASimple:
def __init__(self, mean_flat, components, explained_variance):
self.mean_ = mean_flat
self.components_ = components
self.explained_variance_ = explained_variance

def transform(self, X):
return (X - self.mean_) @ self.components_.T

def inverse_transform(self, b):
return self.mean_ + b @ self.components_

# ---------- ACTIVE SHAPE MODEL ----------
def fit_active_shape_model_with_history(new_mesh, mean_shape, pca, target_points, n_iterations=20, step_size=0.5):
"""
Iterative ASM fitting:
- Align once to target mesh
- Keep correspondences fixed
- Adjust PCA coefficients each iteration
"""
# Align once at start
history = [target_points.copy()]
history.append(mean_shape)
fitted_shape, R_init, t_init, s_init = rigid_procrustes_with_scaling(target_points, mean_shape)

pca_model = PCASimple(mean_shape.flatten(), pca.components_, pca.explained_variance_)
b = np.zeros(pca_model.components_.shape[0]) # start from mean

history.append(fitted_shape.copy())

for _ in range(n_iterations):
# Find target correspondences
target_points = sample_mesh_to_points(new_mesh, fitted_shape)

# Project target to PCA space
target_b = pca_model.transform(target_points.flatten())

# Step toward target in PCA space
b = b + step_size * (target_b - b)

# Clip b to plausible range
stds = np.sqrt(pca_model.explained_variance_)
b = np.clip(b, -5 * stds, 5 * stds)

# Reconstruct new shape
fitted_shape = pca_model.inverse_transform(b).reshape(-1, 3)

history.append(fitted_shape.copy())

return history, b, fitted_shape


This is the result:



My guess is that it's not working well because of the complex shape of the vertebra and probably also because of a displacement of parts. So sometimes there is no overlap between the mesh and the mean shape.

I tried to use multiple or really few points. Also I tried different methods of finding the corresponding point on the new mesh and then fitting to the predicted points instead of aligning to the mesh. Here there was already a problem with the prediction of the points.

As data I used the opensourced verse dataset-01training and dataset-03training. The datasat contains vertebra segmentations in nifti file-format, which needs to be converted to mesh file formats.

Continue reading...