| ## 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() |
