Upload calib_utils.py

calib_utils.py

@@ -0,0 +1,212 @@
+import cv2 as cv
+import numpy as np
+
+pyramid_source = []
+pyramid_target = []
+patch_source_pyramid = []
+JT_source_pyramid = []
+Hinv_source_pyramid = []
+windows = []
+
+
+class Window:
+    def __init__(self, center_x, center_y, window_size):
+        self.center_x = center_x
+        self.center_y = center_y
+        self.window_size = window_size
+        # The displacement vectors.
+        # Important for simulating the internal calculation vector: g and d. Also they record the final results.
+        self.Dx = 0
+        self.Dy = 0
+
+        self.map_x = None
+        self.map_y = None
+        self.generate_map()
+
+    def generate_map(self):
+        epsilon = 0.001
+        start_x = self.center_x - self.window_size//2
+        start_y = self.center_y - self.window_size//2
+        # print(start_x, start_y)
+        # When window_size is odd, we must use this form to enforce the size of map to be window_size.
+        crop_x = np.arange(start_x, start_x+self.window_size - epsilon, 1.0).astype(np.float32).reshape(1, self.window_size)
+        self.map_x = np.repeat(crop_x, self.window_size, axis=0)
+        crop_y = np.arange(start_y, start_y + self.window_size - epsilon, 1.0).astype(np.float32).reshape(self.window_size, 1)
+        self.map_y = np.repeat(crop_y, self.window_size, axis=1)
+
+    def pyrDown(self):
+        """
+        When down-sample the original patch, the corresponding point position should be /2.
+        However the maps' coordinate should not be dimply /2, therefore the maps need regenerate.
+        """
+        self.center_x = self.center_x / 2
+        self.center_y = self.center_y / 2
+        self.Dx = self.Dx / 2
+        self.Dy = self.Dy / 2
+        self.generate_map()
+
+    def pyrUp(self):
+        """
+        When calculating the pyramidal LK and moving to the next (bigger) pyramid, the patch size will be doubled.
+        Thus the corresponding point position should be *2.
+
+        Here we should consider the displacement vector (Dx, Dy), to simulate the equation: g_(L-1) = 2*(g_L + d_L)
+        (d_L calculated in this level iteration and g_L is inherited from the former level iteration, both stored in
+        displacement vector)
+        """
+        self.center_x = self.center_x * 2
+        self.center_y = self.center_y * 2
+        self.Dx = self.Dx * 2
+        self.Dy = self.Dy * 2
+        self.generate_map()
+
+    def move(self, delta_x, delta_y):
+        self.Dx += delta_x
+        self.Dy += delta_y
+
+    def crop(self, img):
+        # Notice!!: map_column calculated from x, while map_row calculated from y.
+        #           Which contradict to the matrix index.
+        patch = cv.remap(img, self.map_x + self.Dx,
+                         self.map_y + self.Dy, cv.INTER_LINEAR)
+        return patch
+
+
+def generate_weight(patch_size):
+    """
+    Generate the weight matrix
+    :param patch_size: (Int) The patch_size
+    :return: The weight map (patch_size * patch_size * 1).
+    """
+    center = [patch_size // 2, patch_size // 2]
+    sigma_x = sigma_y = patch_size // 2
+    maps = np.fromfunction(lambda x, y: ((x - center[0])/sigma_x) ** 2 +
+                                        ((y - center[1])/sigma_y) ** 2,
+                           (patch_size, patch_size),
+                           dtype=int)
+    return np.expand_dims(np.exp(maps/-2.0), -1)
+
+
+def craft_pyramid(image, level, pyramid_container):
+    pyramid_container.clear()
+    pyramid_container.append(image)
+    for i in range(level - 1):
+        image = cv.pyrDown(image)
+        pyramid_container.append(image)
+
+
+def lk_track(face_source, face_target, landmarks_source, window_size, pyramid_level):
+    # Create the image pyramid for both source and target.
+    craft_pyramid(face_source, pyramid_level, pyramid_source)
+    craft_pyramid(face_target, pyramid_level, pyramid_target)
+
+    # Generate the weight map
+    weight_map = generate_weight(window_size)
+
+    # Create windows for cropping patches.
+    windows.clear()
+    for landmark in landmarks_source:
+        x, y = landmark
+        # windows.append(Window(x, y, patch_size, face_source.shape[0], face_source.shape[0]))
+        windows.append(Window(x, y, window_size))
+
+    # Initialize the patches of both the source.
+    # Notice that here both using the same window, i.e., d = 0.
+    # Afterwards, patch_target will be changed while patch_source will fixed.
+    patch_source_pyramid.clear()
+    JT_source_pyramid.clear()
+    Hinv_source_pyramid.clear()
+
+    for level in range(pyramid_level):
+        patch_source = []
+        for window in windows:
+            patch_source.append(window.crop(pyramid_source[level]))
+            if level < pyramid_level - 1:
+                window.pyrDown()
+
+        # Calculate the Jacobian and Hessen matrix of patch_source
+        JT_source = []
+        Hinv_source = []
+        for patch in patch_source:
+            """
+            # cv.Sobel(_, _, x, y, ...), x indicating the horizontal, 
+            # while it's in fact the y axis, for the y is the column.
+            # horizontal means increase at column.
+            """
+            gradient_x = cv.Sobel(patch, cv.CV_64F, 1, 0, ksize=3)
+            gradient_y = cv.Sobel(patch, cv.CV_64F, 0, 1, ksize=3)
+            gradient_x_w = gradient_x * weight_map
+            gradient_y_w = gradient_y * weight_map
+
+            J_x = np.reshape(gradient_x, (-1, 1))
+            J_y = np.reshape(gradient_y, (-1, 1))
+            J_x_w = np.reshape(gradient_x_w, (-1, 1))
+            J_y_w = np.reshape(gradient_y_w, (-1, 1))
+
+            J = np.concatenate((J_x, J_y), axis=1)
+            J_w = np.concatenate((J_x_w, J_y_w), axis=1)
+            JT_w = np.transpose(J_w)
+            H = np.matmul(JT_w, J)
+            Hinv = np.linalg.inv(H)
+            # Noticed that we only collect the weighted JT here.
+            JT_source.append(JT_w)
+            Hinv_source.append(Hinv)
+
+        # Collect all the pre-processed data in each level.
+        patch_source_pyramid.append(patch_source)
+        JT_source_pyramid.append(JT_source)
+        Hinv_source_pyramid.append(Hinv_source)
+    #
+    # """
+    # Sequential Execution
+    # """
+    max_iter_step = 15
+    for level in range(pyramid_level-1, -1, -1):
+        epsilon_der1 = 1.0 + level
+        for patch_s, window, JT, Hinv in zip(patch_source_pyramid[level], windows, JT_source_pyramid[level], Hinv_source_pyramid[level]):
+            count = 1
+            while True:
+                # Patch of target. which will move in each iteration.
+                patch_t = window.crop(pyramid_target[level])
+                # Calculate the residual
+                r = patch_t - patch_s
+                r = np.reshape(r, (-1, 1))
+                der1 = np.matmul(JT, r)
+                der1_norm = np.linalg.norm(der1)
+                delta = - np.matmul(Hinv, der1)
+                if der1_norm < epsilon_der1 or count > max_iter_step:
+                    if level != 0:
+                        # When reach the final level, stop the up-sample.
+                        window.pyrUp()
+                    break
+                else:
+                    window.move(delta[0][0], delta[1][0])
+                    count += 1
+    predictions = []
+    for window in windows:  # type: Window
+        predictions.append([window.center_x + window.Dx, window.center_y + window.Dy])
+    return np.array(predictions)
+
+
+def track_bidirectional(faces, locations):
+    patch_size = 15
+    frames_num = len(faces)
+    pyramid_level = 4
+
+    forward_pts = [locations[0].copy()]
+    for i in range(1, frames_num):
+        feature_old = faces[i-1] / 255.0
+        feature_new = faces[i] / 255.0
+        location_old = forward_pts[i - 1]
+        forward_pt = lk_track(feature_old, feature_new, location_old, patch_size, pyramid_level)
+        forward_pts.append(forward_pt)
+
+    feedback_pts = [None] * (frames_num - 1) + [forward_pts[-1].copy()]
+    for i in range(frames_num - 2, -1, -1):
+        feature_old = faces[i+1] / 255.0
+        feature_new = faces[i] / 255.0
+        location_old = feedback_pts[i - 1]
+        feedback_pt = lk_track(feature_old, feature_new, location_old, patch_size, pyramid_level)
+        feedback_pts[i] = feedback_pt
+
+    return forward_pts, feedback_pts