| import numpy as np |
| from scipy.spatial.distance import pdist, cdist, squareform |
| from sklearn.metrics import pairwise_distances |
|
|
| class ThinPlateSplines: |
| def __init__(self, ctrl_pts: np.ndarray, target_pts: np.ndarray, reg=0.0): |
| """ |
| |
| :param ctrl_pts: [N, d] tensor of control d-dimensional points |
| :param target_pts: [N, d] tensor of target d-dimensional points |
| :param reg: regularization coefficient |
| """ |
| self.__ctrl_pts = ctrl_pts |
| self.__target_pts = target_pts |
| self.__reg = reg |
| self.__num_ctrl_pts = ctrl_pts.shape[0] |
| self.__dim = ctrl_pts.shape[1] |
|
|
| self.__K = None |
| self.__compute_coeffs() |
| self.__aff_params = self.__coeffs[self.__num_ctrl_pts:, ...] |
| self.__non_aff_paramms = self.__coeffs[:self.__num_ctrl_pts, ...] |
|
|
| def __compute_coeffs(self): |
| target_pts_aug = np.vstack([self.__target_pts, |
| np.zeros([self.__dim + 1, self.__dim])]).astype(self.__target_pts.dtype) |
|
|
| T_i = np.linalg.inv(self.__make_T()).astype(self.__target_pts.dtype) |
| self.__coeffs = np.matmul(T_i, target_pts_aug).astype(self.__target_pts.dtype) |
|
|
| def __make_T(self): |
| |
| |
| P = np.hstack([np.ones([self.__num_ctrl_pts, 1], dtype=np.float), self.__ctrl_pts]) |
| zeros = np.zeros([self.__dim + 1, self.__dim + 1], dtype=np.float) |
| self.__K = self.__U_dist(self.__ctrl_pts) |
| alfa = np.mean(self.__K) |
|
|
| self.__K = self.__K + np.ones_like(self.__K) * np.power(alfa, 2) * self.__reg |
|
|
| top = np.hstack([P, self.__K]) |
| bottom = np.hstack([P.transpose(), zeros]) |
|
|
| return np.vstack([top, bottom]) |
|
|
| def __U_dist(self, ctrl_pts, int_pts=None): |
| dist = pairwise_distances(ctrl_pts, int_pts) |
|
|
| if ctrl_pts.shape[-1] == 2: |
| u_dist = np.square(dist) * np.log(dist + 1e-6) |
| else: |
| u_dist = np.sqrt(dist) |
|
|
| return u_dist |
|
|
| def __lift_pts(self, int_pts: np.ndarray, num_pts): |
| |
| |
| |
|
|
| int_pts_lift = np.hstack([self.__U_dist(self.__ctrl_pts, int_pts), |
| np.ones([num_pts, 1], dtype=np.float), |
| int_pts]) |
| return int_pts_lift |
|
|
| def _get_coefficients(self): |
| return self.__coeffs |
|
|
| def interpolate(self, int_points): |
| """ |
| |
| :param int_points: [K, d] flattened d-points of a mesh |
| :return: |
| """ |
| num_pts = int_points.shape[0] |
| int_points_lift = self.__lift_pts(int_points, num_pts) |
| return np.dot(int_points_lift, self.__coeffs) |
|
|
| @property |
| def bending_energy(self): |
| aux = tf.matmul(self.__non_aff_paramms, self.__K, transpose_a=True) |
| return tf.matmul(aux, self.__non_aff_paramms) |
|
|
