diff options
Diffstat (limited to 'tools/3D-Reconstruction/MotionEST/SearchSmooth.py')
-rw-r--r-- | tools/3D-Reconstruction/MotionEST/SearchSmooth.py | 221 |
1 files changed, 221 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..2dc6771ee --- /dev/null +++ b/tools/3D-Reconstruction/MotionEST/SearchSmooth.py @@ -0,0 +1,221 @@ +## Copyright (c) 2020 The WebM project authors. All Rights Reserved. +## +## Use of this source code is governed by a BSD-style license +## that can be found in the LICENSE file in the root of the source +## tree. An additional intellectual property rights grant can be found +## in the file PATENTS. All contributing project authors may +## be found in the AUTHORS file in the root of the source tree. +## + +# 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 |