##  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-8import numpy as npimport numpy.linalg as LAfrom Util import MSEfrom 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