Add Search Smooth Models[Adapt/Fix]

Change-Id: Ia88d16a14b0525d880ac17a133700431949ece31

1 files changed, 213 insertions, 0 deletions

diff --git a/tools/3D-Reconstruction/MotionEST/SearchSmooth.py b/tools/3D-Reconstruction/MotionEST/SearchSmooth.py
new file mode 100644
index 000000000..e0b7fb638
--- /dev/null
+++ b/tools/3D-Reconstruction/MotionEST/SearchSmooth.py
@@ -0,0 +1,213 @@
+#!/usr/bin/env python
+# coding: utf-8
+import numpy as np
+import numpy.linalg as LA
+from Util import MSE
+from MotionEST import MotionEST
+"""Search & Smooth Model with Adapt Weights"""
+
+
+class SearchSmoothAdapt(MotionEST):
+  """
+    Constructor:
+        cur_f: current frame
+        ref_f: reference frame
+        blk_sz: block size
+        wnd_size: search window size
+        beta: neigbor loss weight
+        max_iter: maximum number of iterations
+        metric: metric to compare the blocks distrotion
+    """
+
+  def __init__(self, cur_f, ref_f, blk_size, search, max_iter=100):
+    self.search = search
+    self.max_iter = max_iter
+    super(SearchSmoothAdapt, self).__init__(cur_f, ref_f, blk_size)
+
+  """
+    get local diffiencial of refernce
+    """
+
+  def getRefLocalDiff(self, mvs):
+    m, n = self.num_row, self.num_col
+    localDiff = [[] for _ in xrange(m)]
+    blk_sz = self.blk_sz
+    for r in xrange(m):
+      for c in xrange(n):
+        I_row = 0
+        I_col = 0
+        #get ssd surface
+        count = 0
+        center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz,
+                              c * blk_sz:(c + 1) * blk_sz, 0]
+        ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz)
+        tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz)
+        target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0]
+        for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}:
+          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
+            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
+            I_row += np.sum(np.abs(nb - center)) - np.sum(
+                np.abs(target - center))
+            count += 1
+        I_row //= (count * blk_sz * blk_sz)
+        count = 0
+        for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}:
+          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
+            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
+            I_col += np.sum(np.abs(nb - center)) - np.sum(
+                np.abs(target - center))
+            count += 1
+        I_col //= (count * blk_sz * blk_sz)
+        localDiff[r].append(
+            np.array([[I_row * I_row, I_row * I_col],
+                      [I_col * I_row, I_col * I_col]]))
+    return localDiff
+
+  """
+    add smooth constraint
+    """
+
+  def smooth(self, uvs, mvs):
+    sm_uvs = np.zeros(uvs.shape)
+    blk_sz = self.blk_sz
+    for r in xrange(self.num_row):
+      for c in xrange(self.num_col):
+        nb_uv = np.array([0.0, 0.0])
+        for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}:
+          if 0 <= i < self.num_row and 0 <= j < self.num_col:
+            nb_uv += uvs[i, j] / 6.0
+          else:
+            nb_uv += uvs[r, c] / 6.0
+        for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1),
+                     (r + 1, c + 1)}:
+          if 0 <= i < self.num_row and 0 <= j < self.num_col:
+            nb_uv += uvs[i, j] / 12.0
+          else:
+            nb_uv += uvs[r, c] / 12.0
+        ssd_nb = self.block_dist(r, c, self.blk_sz * nb_uv)
+        mv = mvs[r, c]
+        ssd_mv = self.block_dist(r, c, mv)
+        alpha = (ssd_nb - ssd_mv) / (ssd_mv + 1e-6)
+        M = alpha * self.localDiff[r][c]
+        P = M + np.identity(2)
+        inv_P = LA.inv(P)
+        sm_uvs[r, c] = np.dot(inv_P, nb_uv) + np.dot(
+            np.matmul(inv_P, M), mv / blk_sz)
+    return sm_uvs
+
+  def block_matching(self):
+    self.search.motion_field_estimation()
+
+  def motion_field_estimation(self):
+    self.localDiff = self.getRefLocalDiff(self.search.mf)
+    #get matching results
+    mvs = self.search.mf
+    #add smoothness constraint
+    uvs = mvs / self.blk_sz
+    for _ in xrange(self.max_iter):
+      uvs = self.smooth(uvs, mvs)
+    self.mf = uvs * self.blk_sz
+
+
+"""Search & Smooth Model with Fixed Weights"""
+
+
+class SearchSmoothFix(MotionEST):
+  """
+    Constructor:
+        cur_f: current frame
+        ref_f: reference frame
+        blk_sz: block size
+        wnd_size: search window size
+        beta: neigbor loss weight
+        max_iter: maximum number of iterations
+        metric: metric to compare the blocks distrotion
+    """
+
+  def __init__(self, cur_f, ref_f, blk_size, search, beta, max_iter=100):
+    self.search = search
+    self.max_iter = max_iter
+    self.beta = beta
+    super(SearchSmoothFix, self).__init__(cur_f, ref_f, blk_size)
+
+  """
+    get local diffiencial of refernce
+    """
+
+  def getRefLocalDiff(self, mvs):
+    m, n = self.num_row, self.num_col
+    localDiff = [[] for _ in xrange(m)]
+    blk_sz = self.blk_sz
+    for r in xrange(m):
+      for c in xrange(n):
+        I_row = 0
+        I_col = 0
+        #get ssd surface
+        count = 0
+        center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz,
+                              c * blk_sz:(c + 1) * blk_sz, 0]
+        ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz)
+        tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz)
+        target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0]
+        for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}:
+          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
+            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
+            I_row += np.sum(np.abs(nb - center)) - np.sum(
+                np.abs(target - center))
+            count += 1
+        I_row //= (count * blk_sz * blk_sz)
+        count = 0
+        for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}:
+          if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
+            nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
+            I_col += np.sum(np.abs(nb - center)) - np.sum(
+                np.abs(target - center))
+            count += 1
+        I_col //= (count * blk_sz * blk_sz)
+        localDiff[r].append(
+            np.array([[I_row * I_row, I_row * I_col],
+                      [I_col * I_row, I_col * I_col]]))
+    return localDiff
+
+  """
+    add smooth constraint
+    """
+
+  def smooth(self, uvs, mvs):
+    sm_uvs = np.zeros(uvs.shape)
+    blk_sz = self.blk_sz
+    for r in xrange(self.num_row):
+      for c in xrange(self.num_col):
+        nb_uv = np.array([0.0, 0.0])
+        for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}:
+          if 0 <= i < self.num_row and 0 <= j < self.num_col:
+            nb_uv += uvs[i, j] / 6.0
+          else:
+            nb_uv += uvs[r, c] / 6.0
+        for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1),
+                     (r + 1, c + 1)}:
+          if 0 <= i < self.num_row and 0 <= j < self.num_col:
+            nb_uv += uvs[i, j] / 12.0
+          else:
+            nb_uv += uvs[r, c] / 12.0
+        mv = mvs[r, c] / blk_sz
+        M = self.localDiff[r][c]
+        P = M + self.beta * np.identity(2)
+        inv_P = LA.inv(P)
+        sm_uvs[r, c] = np.dot(inv_P, self.beta * nb_uv) + np.dot(
+            np.matmul(inv_P, M), mv)
+    return sm_uvs
+
+  def block_matching(self):
+    self.search.motion_field_estimation()
+
+  def motion_field_estimation(self):
+    #get local structure
+    self.localDiff = self.getRefLocalDiff(self.search.mf)
+    #get matching results
+    mvs = self.search.mf
+    #add smoothness constraint
+    uvs = mvs / self.blk_sz
+    for _ in xrange(self.max_iter):
+      uvs = self.smooth(uvs, mvs)
+    self.mf = uvs * self.blk_sz