ref: 1d66ec91dac54dac4809f9c56e0b6572c2675dc7
parent: 993173034015fb490a0750895df0b53e3bf41ad7
parent: e436930c39acaf71fed5ca830dcac6878de95b9d
author: Dan Zhu <zxdan@google.com>
date: Fri Aug 23 15:01:00 EDT 2019
Merge "Add Search Smooth Models[Adapt/Fix]"
--- /dev/null
+++ b/tools/3D-Reconstruction/MotionEST/SearchSmooth.py
@@ -1,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