| ## 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 |
| """Exhaust Search:""" |
| |
| |
| class Exhaust(MotionEST): |
| """ |
| Constructor: |
| cur_f: current frame |
| ref_f: reference frame |
| blk_sz: block size |
| wnd_size: search window size |
| metric: metric to compare the blocks distrotion |
| """ |
| |
| def __init__(self, cur_f, ref_f, blk_size, wnd_size, metric=MSE): |
| self.name = 'exhaust' |
| self.wnd_sz = wnd_size |
| self.metric = metric |
| super(Exhaust, self).__init__(cur_f, ref_f, blk_size) |
| |
| """ |
| search method: |
| cur_r: start row |
| cur_c: start column |
| """ |
| |
| def search(self, cur_r, cur_c): |
| min_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) |
| cur_x = cur_c * self.blk_sz |
| cur_y = cur_r * self.blk_sz |
| ref_x = cur_x |
| ref_y = cur_y |
| #search all validate positions and select the one with minimum distortion |
| for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): |
| for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): |
| if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: |
| loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], |
| self.metric) |
| if loss < min_loss: |
| min_loss = loss |
| ref_x = x |
| ref_y = y |
| return ref_x, ref_y |
| |
| def motion_field_estimation(self): |
| for i in xrange(self.num_row): |
| for j in xrange(self.num_col): |
| ref_x, ref_y = self.search(i, j) |
| self.mf[i, j] = np.array( |
| [ref_y - i * self.blk_sz, ref_x - j * self.blk_sz]) |
| |
| |
| """Exhaust with Neighbor Constraint""" |
| |
| |
| class ExhaustNeighbor(MotionEST): |
| """ |
| Constructor: |
| cur_f: current frame |
| ref_f: reference frame |
| blk_sz: block size |
| wnd_size: search window size |
| beta: neigbor loss weight |
| metric: metric to compare the blocks distrotion |
| """ |
| |
| def __init__(self, cur_f, ref_f, blk_size, wnd_size, beta, metric=MSE): |
| self.name = 'exhaust + neighbor' |
| self.wnd_sz = wnd_size |
| self.beta = beta |
| self.metric = metric |
| super(ExhaustNeighbor, self).__init__(cur_f, ref_f, blk_size) |
| self.assign = np.zeros((self.num_row, self.num_col), dtype=np.bool) |
| |
| """ |
| estimate neighbor loss: |
| cur_r: current row |
| cur_c: current column |
| mv: current motion vector |
| """ |
| |
| def neighborLoss(self, cur_r, cur_c, mv): |
| loss = 0 |
| #accumulate difference between current block's motion vector with neighbors' |
| for i, j in {(-1, 0), (1, 0), (0, 1), (0, -1)}: |
| nb_r = cur_r + i |
| nb_c = cur_c + j |
| if 0 <= nb_r < self.num_row and 0 <= nb_c < self.num_col and self.assign[ |
| nb_r, nb_c]: |
| loss += LA.norm(mv - self.mf[nb_r, nb_c]) |
| return loss |
| |
| """ |
| search method: |
| cur_r: start row |
| cur_c: start column |
| """ |
| |
| def search(self, cur_r, cur_c): |
| dist_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) |
| nb_loss = self.neighborLoss(cur_r, cur_c, np.array([0, 0])) |
| min_loss = dist_loss + self.beta * nb_loss |
| cur_x = cur_c * self.blk_sz |
| cur_y = cur_r * self.blk_sz |
| ref_x = cur_x |
| ref_y = cur_y |
| #search all validate positions and select the one with minimum distortion |
| # as well as weighted neighbor loss |
| for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): |
| for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): |
| if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: |
| dist_loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], |
| self.metric) |
| nb_loss = self.neighborLoss(cur_r, cur_c, [y - cur_y, x - cur_x]) |
| loss = dist_loss + self.beta * nb_loss |
| if loss < min_loss: |
| min_loss = loss |
| ref_x = x |
| ref_y = y |
| return ref_x, ref_y |
| |
| def motion_field_estimation(self): |
| for i in xrange(self.num_row): |
| for j in xrange(self.num_col): |
| ref_x, ref_y = self.search(i, j) |
| self.mf[i, j] = np.array( |
| [ref_y - i * self.blk_sz, ref_x - j * self.blk_sz]) |
| self.assign[i, j] = True |
| |
| |
| """Exhaust with Neighbor Constraint and Feature Score""" |
| |
| |
| class ExhaustNeighborFeatureScore(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, |
| wnd_size, |
| beta=1, |
| max_iter=100, |
| metric=MSE): |
| self.name = 'exhaust + neighbor+feature score' |
| self.wnd_sz = wnd_size |
| self.beta = beta |
| self.metric = metric |
| self.max_iter = max_iter |
| super(ExhaustNeighborFeatureScore, self).__init__(cur_f, ref_f, blk_size) |
| self.fs = self.getFeatureScore() |
| |
| """ |
| get feature score of each block |
| """ |
| |
| def getFeatureScore(self): |
| fs = np.zeros((self.num_row, self.num_col)) |
| for r in xrange(self.num_row): |
| for c in xrange(self.num_col): |
| IxIx = 0 |
| IyIy = 0 |
| IxIy = 0 |
| #get ssd surface |
| for x in xrange(self.blk_sz - 1): |
| for y in xrange(self.blk_sz - 1): |
| ox = c * self.blk_sz + x |
| oy = r * self.blk_sz + y |
| Ix = self.cur_yuv[oy, ox + 1, 0] - self.cur_yuv[oy, ox, 0] |
| Iy = self.cur_yuv[oy + 1, ox, 0] - self.cur_yuv[oy, ox, 0] |
| IxIx += Ix * Ix |
| IyIy += Iy * Iy |
| IxIy += Ix * Iy |
| #get maximum and minimum eigenvalues |
| lambda_max = 0.5 * ((IxIx + IyIy) + np.sqrt(4 * IxIy * IxIy + |
| (IxIx - IyIy)**2)) |
| lambda_min = 0.5 * ((IxIx + IyIy) - np.sqrt(4 * IxIy * IxIy + |
| (IxIx - IyIy)**2)) |
| fs[r, c] = lambda_max * lambda_min / (1e-6 + lambda_max + lambda_min) |
| if fs[r, c] < 0: |
| fs[r, c] = 0 |
| return fs |
| |
| """ |
| do exhaust search |
| """ |
| |
| def search(self, cur_r, cur_c): |
| min_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) |
| cur_x = cur_c * self.blk_sz |
| cur_y = cur_r * self.blk_sz |
| ref_x = cur_x |
| ref_y = cur_y |
| #search all validate positions and select the one with minimum distortion |
| for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): |
| for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): |
| if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: |
| loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], |
| self.metric) |
| if loss < min_loss: |
| min_loss = loss |
| ref_x = x |
| ref_y = y |
| return ref_x, ref_y |
| |
| """ |
| add smooth constraint |
| """ |
| |
| def smooth(self, uvs, mvs): |
| sm_uvs = np.zeros(uvs.shape) |
| for r in xrange(self.num_row): |
| for c in xrange(self.num_col): |
| avg_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: |
| avg_uv += uvs[i, j] / 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: |
| avg_uv += uvs[i, j] / 12.0 |
| sm_uvs[r, c] = (self.fs[r, c] * mvs[r, c] + self.beta * avg_uv) / ( |
| self.beta + self.fs[r, c]) |
| return sm_uvs |
| |
| def motion_field_estimation(self): |
| #get matching results |
| mvs = np.zeros(self.mf.shape) |
| for r in xrange(self.num_row): |
| for c in xrange(self.num_col): |
| ref_x, ref_y = self.search(r, c) |
| mvs[r, c] = np.array([ref_y - r * self.blk_sz, ref_x - c * self.blk_sz]) |
| #add smoothness constraint |
| uvs = np.zeros(self.mf.shape) |
| for _ in xrange(self.max_iter): |
| uvs = self.smooth(uvs, mvs) |
| self.mf = uvs |
