| ## 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 | |
| import matplotlib.pyplot as plt | |
| from 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() |