##  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 LAimport matplotlib.pyplot as pltfrom Util import drawMF, MSE"""The Base Class of Estimators"""  class MotionEST(object):  """    constructor:        cur_f: current frame        ref_f: reference frame        blk_sz: block size    """   def __init__(self, cur_f, ref_f, blk_sz):    self.cur_f = cur_f    self.ref_f = ref_f    self.blk_sz = blk_sz    #convert RGB to YUV    self.cur_yuv = np.array(self.cur_f.convert('YCbCr'), dtype=np.int)    self.ref_yuv = np.array(self.ref_f.convert('YCbCr'), dtype=np.int)    #frame size    self.width = self.cur_f.size[0]    self.height = self.cur_f.size[1]    #motion field size    self.num_row = self.height // self.blk_sz    self.num_col = self.width // self.blk_sz    #initialize motion field    self.mf = np.zeros((self.num_row, self.num_col, 2))   """estimation function Override by child classes"""   def motion_field_estimation(self):    pass   """    distortion of a block:      cur_r: current row      cur_c: current column      mv: motion vector      metric: distortion metric  """   def block_dist(self, cur_r, cur_c, mv, metric=MSE):    cur_x = cur_c * self.blk_sz    cur_y = cur_r * self.blk_sz    h = min(self.blk_sz, self.height - cur_y)    w = min(self.blk_sz, self.width - cur_x)    cur_blk = self.cur_yuv[cur_y:cur_y + h, cur_x:cur_x + w, :]    ref_x = int(cur_x + mv[1])    ref_y = int(cur_y + mv[0])    if 0 <= ref_x < self.width - w and 0 <= ref_y < self.height - h:      ref_blk = self.ref_yuv[ref_y:ref_y + h, ref_x:ref_x + w, :]    else:      ref_blk = np.zeros((h, w, 3))    return metric(cur_blk, ref_blk)   """    distortion of motion field  """   def distortion(self, mask=None, metric=MSE):    loss = 0    count = 0    for i in xrange(self.num_row):      for j in xrange(self.num_col):        if mask is not None and mask[i, j]:          continue        loss += self.block_dist(i, j, self.mf[i, j], metric)        count += 1    return loss / count   """evaluation compare the difference with ground truth"""   def motion_field_evaluation(self, ground_truth):    loss = 0    count = 0    gt = ground_truth.mf    mask = ground_truth.mask    for i in xrange(self.num_row):      for j in xrange(self.num_col):        if mask is not None and mask[i][j]:          continue        loss += LA.norm(gt[i, j] - self.mf[i, j])        count += 1    return loss / count   """render the motion field"""   def show(self, ground_truth=None, size=10):    cur_mf = drawMF(self.cur_f, self.blk_sz, self.mf)    if ground_truth is None:      n_row = 1    else:      gt_mf = drawMF(self.cur_f, self.blk_sz, ground_truth)      n_row = 2    plt.figure(figsize=(n_row * size, size * self.height / self.width))    plt.subplot(1, n_row, 1)    plt.imshow(cur_mf)    plt.title('Estimated Motion Field')    if ground_truth is not None:      plt.subplot(1, n_row, 2)      plt.imshow(gt_mf)      plt.title('Ground Truth')    plt.tight_layout()    plt.show()