Merge pull request #9 from Samleo8/motionDistort

Motion distort

genFakeData.py

@@ -1,6 +1,6 @@
 import numpy as np
 import matplotlib.pyplot as plt
-from utils import getRotationMatrix
+from utils import getRotationMatrix, invert_transform
 from parseData import *
 
 from parseData import RANGE_RESOLUTION_CART_M
@@ -40,23 +40,23 @@ def plotFakeFeatures(srcCoord,
                     color='blue',
                     marker='.',
                     alpha=alpha,
-                    label=f'Features 0{title_append}')
+                    label=f'Instantaneous Radar Scan{title_append}')
 
     if targetCoord is not None:
         plt.scatter(targetCoord[:, 0],
                     targetCoord[:, 1],
                     color='red',
                     marker='+',
                     alpha=alpha,
-                    label=f'Features 1{title_append}')
+                    label=f'Scan with Distortion{title_append}')
 
     if targetCoord2 is not None:
         plt.scatter(targetCoord2[:, 0],
                     targetCoord2[:, 1],
                     color='green',
                     marker='x',
                     alpha=alpha,
-                    label=f'Features 2{title_append}')
+                    label=f'Original Points{title_append}')
 
     if plotDisplace:
         for i in range(targetCoord.shape[0]):
@@ -116,44 +116,44 @@ def convertPolarPointsToCartesian(points):
     y = np.expand_dims(ranges * np.sin(angles), axis = 1)
     return np.hstack((x, y))
 
-def generateFakeCorrespondencesPolar(srcCoord=None,
+def generateFakeCorrespondencesPolar(currentFrame=None,
                                     n_points=100,
                                     theta_max_deg=20,
                                     max_translation_m=3):
     '''
     @brief Generate fake correspondences with transform, randomly generated from max range and degree
-    @param[in] srcCoord Source coordinate to transform from. If none, will randomly generate features
-    @param[in] n_points Number of points to generate, only applies if srcCoord = None
+    @param[in] currentFrame Source coordinate to transform from. If none, will randomly generate features
+    @param[in] n_points Number of points to generate, only applies if currentFrame = None
     @param[in] theta_max_deg Maximum degree of rotation
     @param[in] max_range_m Maximum range (for translation) in meters
 
-    @return srcCoord Generated or passed in srcCoord
+    @return currentFrame Generated or passed in currentFrame
     @return targetCoord Corresponding targetCoord generated using (theta_deg, h)
     @return theta_deg Theta component of transform
     @return h Translation component of transform
     '''
 
-    if srcCoord is None:
+    if currentFrame is None:
         print("Generating fake features..")
         max_range_m = max_translation_m * 3
-        srcCoord = generateFakeFeaturesPolar(n_points, max_range_m)
-        #print(srcCoord.shape)
-        srcCoord = convertPolarPointsToCartesian(srcCoord)
+        currentFrame = generateFakeFeaturesPolar(n_points, max_range_m)
+        #print(currentFrame.shape)
+        currentFrame = convertPolarPointsToCartesian(currentFrame)
     else:
-        n_points = srcCoord.shape[0]
+        n_points = currentFrame.shape[0]
 
     theta_deg = np.random.random() * theta_max_deg
     R = getRotationMatrix(theta_deg, degrees=True)
     h = generateTranslationVector(max_translation_m)
-    #print(srcCoord.shape)
-    targetCoord = transformCoords(srcCoord, R, h)
+    #h = np.array([[0], [0]])
+    groundTruth = transformCoords(currentFrame, R, h)
 
-    return srcCoord, targetCoord, theta_deg, h
+    return groundTruth, currentFrame, theta_deg, h
 
 def distort(coords, velocity, frequency, h):
 
     coords_norm = coords - h.flatten() # N x 2
-    angles = np.arctan2(coords_norm[:, 1], -coords_norm[:, 0]) # - y to follow clockwise convention
+    angles = np.arctan2(coords_norm[:, 1], -coords_norm[:, 0]) # - x to follow clockwise convention
     period = 1 / frequency
     times = angles / (2 * np.pi) * period
     #print(angles) # lesson: need to use arctan2 wisely, it wraps [-pi, pi]
@@ -166,15 +166,14 @@ def distort(coords, velocity, frequency, h):
     #print(displacement)
     #print(displacement)
     dx = displacement[0, :]
-    print(dx)
     dy = displacement[1, :]
     dtheta = displacement[2, :] / 180 * np.pi
     c = np.cos(dtheta)
     s = np.sin(dtheta)
     ones = np.ones(times.shape)
     zeros = np.zeros(times.shape)
-    distortion = np.array([[c, s, -s* dy - c*dx],
-                           [-s,  c, -c*dy + s*dx],
+    distortion = np.array([[ c, s, -s*dy - c*dx],
+                           [-s, c, -c*dy + s*dx],
                            [zeros, zeros, ones]]) # 3 x 3 x N, need to invert?
     distorted = np.transpose(distortion, axes = (2, 0, 1)) @ np.expand_dims(coords, axis = 2) # N x 3 x 1
     distorted = distorted[:, :2, 0]

getTransformKLT.py

@@ -468,15 +468,16 @@ def getTrackedPointsKLT(
             good_old, good_new = rejectOutliers(good_old, good_new)
 
             # Obtain transforms
-            R, h = calculateTransformDxDth(good_old, good_new)
-            # R, h = calculateTransformSVD(good_old, good_new)
+            #R, h = calculateTransformDxDth(good_old, good_new)
+            R, h = calculateTransformSVD(good_old, good_new)
             # print(h)
             # h[0] += 0
             # for i in range(good_old.shape[0]):
             #     plotFakeFeatures(good_old[i:i+1,:], (R @ good_new[i:i+1,:].T + h).T, show= True)
-            # transformed_pts = (R @ good_new.T + h).T
+            transformed_pts = (R @ good_new.T + h).T
             # print(f"RMSE = {np.sum(np.square(good_old - transformed_pts))}")
-            # plotFakeFeatures(good_old, transformed_pts, good_new, show = True)
+            #plotFakeFeatures(good_old, good_new, show = True)
+            plotFakeFeatures(good_old, transformed_pts, good_new, show = True)
             h *= RANGE_RESOLUTION_CART_M
 
             #R, h = estimateTransformUsingDelats(good_old, good_new)

motionDistortion.py

@@ -1,5 +1,6 @@
 import numpy as np
 import scipy as sp
+from utils import *
 #sp.linalg.sqrtm
 #sp.lin
 
@@ -29,26 +30,32 @@
 Debug
 
 '''
+
+
 class MotionDistortionSolver():
-    def __init__(self, T_wj0, p_w, p_jt, v_j0, T_wj, lambda_p, lambda_v):
+    def __init__(self, T_wj0, p_w, p_jt, v_j0, T_wj, sigma_p, sigma_v):
         # e_p Parameters
         self.T_wj0 = T_wj0 # Prior transform, T_w,j0
         self.T_wj0_inv = np.linalg.inv(T_wj0)
-        self.p_w = p_w # Estimated point world positions, N x 3
-        self.p_jt = p_jt # Observed points at time t, N x 3
-
+        self.p_w = homogenize(p_w) # Estimated point world positions, N x 3
+        self.p_jt = homogenize(p_jt) # Observed points at time t, N x 3
+        
         # e_v Parameters
         self.v_j_initial = v_j0 # Initial velocity guess (prior velocity/ velocity from SVD solution)
         self.T_wj_initial = T_wj # Initial Transform guess (T from SVD solution)
 
         # Optimization parameters
-        self.lambda_p = lambda_p # Info matrix, point error, lamdba_p
-        self.lambda_v = lambda_v # Info matrix, velocity, lambda_v
-        nv = self.lambda_p.shape[0]
-        np = self.lambda_v.shape[0]
-        self.lambda_total = np.block([[lambda_v, np.zeros((nv, np))],
-                                      [np.zeros((np, nv)), lambda_p]])
-        self.info_sqrt = sp.linalg.sqrtm(np.linalg.inv(self.lambda_total)) # 5 x 5
+        self.sigma_p = sigma_p # Info matrix, point error, lamdba_p
+        self.sigma_v = sigma_v # Info matrix, velocity, sigma_v
+        n_v = self.sigma_v.shape[0]
+        n_p = self.sigma_p.shape[0]
+        self.sigma_total = np.block([[sigma_v, np.zeros((n_v, n_p))],
+                                      [np.zeros((n_p, n_v)), sigma_p]])
+        self.info_sqrt = sp.linalg.sqrtm(np.linalg.inv(self.sigma_total)) # 5 x 5
+        sigma_p_vector = np.tile(np.diag(sigma_p), p_jt.shape[0])
+        sigma_v_vector = np.diag(sigma_v)
+        sigma_vector = np.concatenate((sigma_p_vector, sigma_v_vector))
+        self.info_vector = 1 / sigma_vector
         self.dT = None
         self.T_wj_best = T_wj
         self.v_j_best = v_j0 # might not be good enough a guess, too far from optimal
@@ -57,7 +64,7 @@ def __init__(self, T_wj0, p_w, p_jt, v_j0, T_wj, lambda_p, lambda_v):
         self.total_scan_time = 1 / 4 # assuming 4 Hz
         pass
 
-    def compute_time_deltas(self):
+    def compute_time_deltas(self, points):
         '''
         Get the time deltas for each point. This depends solely on where the
         points are in scan angle. The further away from center, the greater the
@@ -67,34 +74,35 @@ def compute_time_deltas(self):
         idea to re-run this function once an undistorted frame is obtained for
         better estimates.
         '''
-        points = self.undistort()#self.p_jt # provide in N x 3
+        #points = self.undistort()#self.p_jt # provide in N x 3
         x = points[:, 0]
         y = points[:, 1]
-        angles = np.arctan2(-y, x) # scanline starts at positive x axis and moves clockwise
+        angles = np.arctan2(y, -x) # scanline starts at positive x axis and moves clockwise, we take midpoint pi/2 as 0
         dT = self.total_scan_time * angles / (2 * np.pi)
-        dT -= self.total_scan_time / 2 # offset range, as defined in [-scan_time /2 , scan_time/2]
+        #dT -= self.total_scan_time / 2 # offset range, as defined in [-scan_time /2 , scan_time/2]
         self.dT = dT
 
     def undistort(self, v_j):
         '''
         Computes a new set of undistorted observed points, based on the current
         best estimate of v_T, T_wj, dT
         '''
-        displacement = np.expand_dims(v_j, axis = 1) * self.dT # 3 x N
-        assert(displacement.shape = (3,points.shape[0]))
-        theta = displacement[2, :]
-        dx = displacement[0, :]
-        dy = displacement[1, :]
+        v_j_column = np.expand_dims(v_j, axis = 1)
+        displacement = v_j_column * self.dT # 3 x N
+
+        theta = displacement[2, :] # (N,)
+        dx = displacement[0, :] # (N,)
+        dy = displacement[1, :] # (N,)
+        shape = theta.shape
         # Correction matrix for time drift, 3 x 3 x N
         T_j_jt = np.array([[np.cos(theta), -np.sin(theta), dx],
                            [np.sin(theta), np.cos(theta),  dy],
-                           [0,             0,              1]])
-
-        p_jt_col = np.expand_dims(self.p_jt, axis = 1).transpose(axis = (2, 0, 1)) # N x 3 x 1
-        undistorted = T_j_jt.transpose(axis = (2, 0, 1)) @ p_jt_col # N x 3 x 1
+                           [np.zeros(shape), np.zeros(shape), np.ones(shape)]])
+        p_jt_col = np.expand_dims(self.p_jt, axis = 2) # N x 3 x 1
+        undistorted = T_j_jt.transpose((2, 0, 1)) @ p_jt_col # N x 3 x 1
         return undistorted
 
-
+    # TODO: Check if this really should be inverted
     def expected_observed_pts(self, T_wj):
         '''
         Returns the estimated positions of points based on their world location
@@ -103,13 +111,19 @@ def expected_observed_pts(self, T_wj):
         return np.linalg.inv(T_wj) @ self.p_w.T
 
     def error_vector(self, params):
-        theta = params[2]
-        x = params[0]
-        y = params[1]
+        '''
+        Because we are optimizing over rotations, we choose to keep the rotation
+        in a theta form, we have to do matrix exponential in here to convert
+        into the SO(1) form, then augment to the rotation-translation transform
+        '''
+        theta = params[5]
+        x = params[3]
+        y = params[4]
         T = np.array([[np.cos(theta), -np.sin(theta), x],
                       [np.sin(theta),  np.cos(theta), y],
                       [0            ,  0            , 1]])
-        return self.info_sqrt @ self.error(params[:3], T)
+        #return self.info_sqrt @ self.error(params[:3], T)
+        return self.info_vector * self.error(params[:3], T)
 
     def error(self, v_j, T_wj):
         '''
@@ -118,32 +132,41 @@ def error(self, v_j, T_wj):
         '''
         # Compute point error
         undistorted = self.undistort(v_j)
-        expected = self.expected_observed_pts(self, T_wj)
+        #self.compute_time_deltas(undistorted)
+        expected = self.expected_observed_pts(T_wj)
         naive_e_p = expected - np.squeeze(undistorted).T # 3 x N
         # Actual loss is the Cauchy robust loss, defined here:
-        e_p_i = np.log(np.square(naive_e_p[:2, :]) / 2 + 1)
-        e_p = np.sum(e_p_i, axis = 1) # 2 x 1
+        e_p_i = np.log(np.square(naive_e_p[:2, :]) / 2 + 1) # 2 x N
+        e_p = e_p_i.flatten(order='F')
+        #e_p = np.sum(e_p_i, axis = 1) # (2,)
 
         # Compute velocity error
+        # Matrix log operation
         T_j_j1 = self.T_wj0_inv @ T_wj
         dx = T_j_j1[0, 2]
         dy = T_j_j1[1, 2]
         dtheta = np.arctan2(T_j_j1[1, 0], T_j_j1[0, 0])
         v_j_prior = np.array([dx, dy, dtheta]) / self.total_scan_time
-        e_v = (v_j - v_j_prior) * e_p.shape[1] # 3 x 1
-
-        e = np.vstack((e_v, e_p))
+        #print(f"Prior velocity: {v_j_prior}")
+        v_diff = (v_j - v_j_prior)
+        v_diff[2] = normalize_angles(v_diff[2])
+        e_v = v_diff * e_p_i.shape[1] # (3,)
+        # ideally should warp2pi here on theta error
+        e = np.hstack((e_p, e_v))
+        #print(e)
         return e
 
     def jacobian_vector(self, params):
-        theta = params[2]
-        x = params[0]
-        y = params[1]
+        theta = params[5]
+        x = params[3]
+        y = params[4]
         T = np.array([[np.cos(theta), -np.sin(theta), x],
                       [np.sin(theta),  np.cos(theta), y],
                       [0            ,  0            , 1]])
-        return self.info_sqrt @ self.jacobian(params[:3], T)
-
+        velocity = params[:3]
+        #return self.info_sqrt @ self.jacobian(velocity, T)
+        return np.expand_dims(self.info_vector, axis=1) * self.jacobian(velocity, T)
+
     def jacobian(self, v_j, T_wj):
         '''
         Compute the Jacobian. This has two parts, as defined by the RadarSLAM
@@ -154,11 +177,12 @@ def jacobian(self, v_j, T_wj):
                 vx, vy, vtheta
         '''
         undistorted = self.undistort(v_j)
-        expected = self.expected_observed_pts(self, T_wj)
+        expected = self.expected_observed_pts(T_wj)
         input = expected - np.squeeze(undistorted).T # 3 x N
-        cauchy_derivative = input / (np.square(input[:2, ]) / 2 + 1) # 2 x N
+        naive_e_p = input[:2]
+        cauchy_derivative = naive_e_p / (np.square(naive_e_p) / 2 + 1) # 3 x N
 
-        # Compute J_p: derivative of e_p wrt
+        # Compute J_p: derivative of errors wrt the point position
         c0 = self.T_wj0[0, 0]
         s0 = self.T_wj0[1, 0]
         c1 = T_wj[0, 0]
@@ -168,14 +192,16 @@ def jacobian(self, v_j, T_wj):
         pwx = self.p_w[:, 0] # 1 x N
         pwy = self.p_w[:, 1]
         ones = np.ones(pwx.shape)
+        zeros = np.zeros(pwx.shape)
 
         # 2 x 3 x N
         J_p1 = np.array([[-c1 * ones, -s1 * ones, -pwx * s1 + pwy * c1 - c1 * Ty + s1 * Tx],
                         [s1 * ones,  -c1 * ones, -pwx * c1 - pwy * s1 + s1 * Ty + c1 * Tx]])
         J_p1 *= np.expand_dims(cauchy_derivative, axis = 1)
+        J_p1 = np.squeeze(np.vstack(np.split(J_p1, J_p1.shape[2], axis = 2)))
         J_p2 = np.array([[ c0, s0, 0],
                          [-s0, c0, 0],
-                         [0,   0,  1]]) / self.total_scan_time
+                         [0,   0,  1]]) / self.total_scan_time * pwx.shape[0]
         J_p = np.vstack((J_p1, J_p2))
 
         # Compute J_v: derivative of the errors wrt velocity
@@ -184,10 +210,11 @@ def jacobian(self, v_j, T_wj):
         y = points[:, 1]
         displacement = np.expand_dims(v_j, axis = 1) * self.dT # 3 x N
         theta = displacement[2, :]
-        J_v = np.array([[-self.dT, 0, np.sin(theta) * self.dT * x + np.cos(theta) * self.dT * y ],
-                        [0, -self.dT, -np.cos(theta) * self.dT * x + np.sin(theta) * self.dT * y]])
+        zeros = np.zeros(theta.shape)
+        J_v = np.array([[-self.dT, zeros, self.dT * (np.sin(theta) * x + np.cos(theta) * y) ],
+                        [zeros, -self.dT, self.dT * (-np.cos(theta) * x + np.sin(theta) * y)]])
         J_v *= np.expand_dims(cauchy_derivative, axis = 1)
-        J_v = np.sum(J_v, axis = -1) # 3 x 2
+        J_v = np.squeeze(np.vstack(np.split(J_v, J_v.shape[2], axis = 2))) # 2N x 3
         J_v = np.vstack((J_v, np.eye(3) * x.shape[0]))
 
         # J = [J_v, J_p]
@@ -218,13 +245,19 @@ def optimize(self, max_iters = 20):
         pass
 
     def optimize_library(self):
+        '''
+        Optimize using the LM implementation in the scipy library.
+        '''
+        self.compute_time_deltas(self.p_jt)
         # Initialize v, T
         T0 = self.T_wj_initial
         T_params = np.array([T0[0, 2], T0[1, 2], np.arctan2(T0[1, 0], T0[0, 0])])
         initial_guess = np.hstack((self.v_j_initial, T_params))
+        print(f"Initial v guess: {self.v_j_initial}")
+        print(f"Initial T guess: {T_params}")
 
-        result = sp.optimize.least_squares(self.error_vector, initial_guess, jac = self.jacobian_vector, method = 'lm')
-
+        result = sp.optimize.least_squares(self.error_vector, initial_guess, jac = '2-point', method = 'lm')
+        # result = sp.optimize.least_squares(self.error_vector, initial_guess, jac = self.jacobian_vector, method = 'lm')
         # return v, T
         best_params = result.x
         num_evals = result.nfev # number of function evaluations: measure of how quickly we converged
@@ -235,7 +268,9 @@ def optimize_library(self):
                         2 : "ftol termination condition is satisfied",
                         3 : "xtol termination condition is satisfied",
                         4 : "Both ftol and xtol termination conditions are satisfied"}
-        print(f"Final v: {best_params[:3]}, t: {best_params[3:]}")
+        print(f"Final v: {best_params[:3]}")
+        print(f"Final t: {best_params[3:]}")
         print(f"Used {num_evals} evaluations")
+        print(f"Residuals were {result.fun}")
         print(status_dict[status])
         return best_params