-
Notifications
You must be signed in to change notification settings - Fork 111
/
rigid_registration.py
152 lines (120 loc) · 5.55 KB
/
rigid_registration.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
from builtins import super
import numpy as np
import numbers
from .emregistration import EMRegistration
from .utility import is_positive_semi_definite
class RigidRegistration(EMRegistration):
"""
Rigid registration.
Attributes
----------
R: numpy array (semi-positive definite)
DxD rotation matrix. Any well behaved matrix will do,
since the next estimate is a rotation matrix.
t: numpy array
1xD initial translation vector.
s: float (positive)
scaling parameter.
A: numpy array
Utility array used to calculate the rotation matrix.
Defined in Fig. 2 of https://arxiv.org/pdf/0905.2635.pdf.
"""
# Additional parameters used in this class, but not inputs.
# YPY: float
# Denominator value used to update the scale factor.
# Defined in Fig. 2 and Eq. 8 of https://arxiv.org/pdf/0905.2635.pdf.
# X_hat: numpy array
# Centered target point cloud.
# Defined in Fig. 2 of https://arxiv.org/pdf/0905.2635.pdf.
def __init__(self, R=None, t=None, s=None, *args, **kwargs):
super().__init__(*args, **kwargs)
if self.D != 2 and self.D != 3:
raise ValueError(
'Rigid registration only supports 2D or 3D point clouds. Instead got {}.'.format(self.D))
if R is not None and (R.ndim is not 2 or R.shape[0] is not self.D or R.shape[1] is not self.D or not is_positive_semi_definite(R)):
raise ValueError(
'The rotation matrix can only be initialized to {}x{} positive semi definite matrices. Instead got: {}.'.format(self.D, self.D, R))
if t is not None and (t.ndim is not 2 or t.shape[0] is not 1 or t.shape[1] is not self.D):
raise ValueError(
'The translation vector can only be initialized to 1x{} positive semi definite matrices. Instead got: {}.'.format(self.D, t))
if s is not None and (not isinstance(s, numbers.Number) or s <= 0):
raise ValueError(
'The scale factor must be a positive number. Instead got: {}.'.format(s))
self.R = np.eye(self.D) if R is None else R
self.t = np.atleast_2d(np.zeros((1, self.D))) if t is None else t
self.s = 1 if s is None else s
def update_transform(self):
"""
Calculate a new estimate of the rigid transformation.
"""
# target point cloud mean
muX = np.divide(np.sum(self.PX, axis=0),
self.Np)
# source point cloud mean
muY = np.divide(
np.sum(np.dot(np.transpose(self.P), self.Y), axis=0), self.Np)
self.X_hat = self.X - np.tile(muX, (self.N, 1))
# centered source point cloud
Y_hat = self.Y - np.tile(muY, (self.M, 1))
self.YPY = np.dot(np.transpose(self.P1), np.sum(
np.multiply(Y_hat, Y_hat), axis=1))
self.A = np.dot(np.transpose(self.X_hat), np.transpose(self.P))
self.A = np.dot(self.A, Y_hat)
# Singular value decomposition as per lemma 1 of https://arxiv.org/pdf/0905.2635.pdf.
U, _, V = np.linalg.svd(self.A, full_matrices=True)
C = np.ones((self.D, ))
C[self.D-1] = np.linalg.det(np.dot(U, V))
# Calculate the rotation matrix using Eq. 9 of https://arxiv.org/pdf/0905.2635.pdf.
self.R = np.transpose(np.dot(np.dot(U, np.diag(C)), V))
# Update scale and translation using Fig. 2 of https://arxiv.org/pdf/0905.2635.pdf.
self.s = np.trace(np.dot(np.transpose(self.A),
np.transpose(self.R))) / self.YPY
self.t = np.transpose(muX) - self.s * \
np.dot(np.transpose(self.R), np.transpose(muY))
def transform_point_cloud(self, Y=None):
"""
Update a point cloud using the new estimate of the rigid transformation.
Attributes
----------
Y: numpy array
Point cloud to be transformed - use to predict on new set of points.
Best for predicting on new points not used to run initial registration.
If None, self.Y used.
Returns
-------
If Y is None, returns None.
Otherwise, returns the transformed Y.
"""
if Y is None:
self.TY = self.s * np.dot(self.Y, self.R) + self.t
return
else:
return self.s * np.dot(Y, self.R) + self.t
def update_variance(self):
"""
Update the variance of the mixture model using the new estimate of the rigid transformation.
See the update rule for sigma2 in Fig. 2 of of https://arxiv.org/pdf/0905.2635.pdf.
"""
qprev = self.q
trAR = np.trace(np.dot(self.A, self.R))
xPx = np.dot(np.transpose(self.Pt1), np.sum(
np.multiply(self.X_hat, self.X_hat), axis=1))
self.q = (xPx - 2 * self.s * trAR + self.s * self.s * self.YPY) / \
(2 * self.sigma2) + self.D * self.Np/2 * np.log(self.sigma2)
self.diff = np.abs(self.q - qprev)
self.sigma2 = (xPx - self.s * trAR) / (self.Np * self.D)
if self.sigma2 <= 0:
self.sigma2 = self.tolerance / 10
def get_registration_parameters(self):
"""
Return the current estimate of the rigid transformation parameters.
Returns
-------
self.s: float
Current estimate of the scale factor.
self.R: numpy array
Current estimate of the rotation matrix.
self.t: numpy array
Current estimate of the translation vector.
"""
return self.s, self.R, self.t