1##  Copyright (c) 2020 The WebM project authors. All Rights Reserved.
2##
3##  Use of this source code is governed by a BSD-style license
4##  that can be found in the LICENSE file in the root of the source
5##  tree. An additional intellectual property rights grant can be found
6##  in the file PATENTS.  All contributing project authors may
7##  be found in the AUTHORS file in the root of the source tree.
8##
9
10# coding: utf-8
11import numpy as np
12import numpy.linalg as LA
13import matplotlib.pyplot as plt
14from scipy.ndimage import filters
15from PIL import Image, ImageDraw
16
17
18def MSE(blk1, blk2):
19  return np.mean(
20      LA.norm(
21          np.array(blk1, dtype=np.int) - np.array(blk2, dtype=np.int), axis=2))
22
23
24def drawMF(img, blk_sz, mf):
25  img_rgba = img.convert('RGBA')
26  mf_layer = Image.new(mode='RGBA', size=img_rgba.size, color=(0, 0, 0, 0))
27  draw = ImageDraw.Draw(mf_layer)
28  width = img_rgba.size[0]
29  height = img_rgba.size[1]
30  num_row = height // blk_sz
31  num_col = width // blk_sz
32  for i in xrange(num_row):
33    left = (0, i * blk_sz)
34    right = (width, i * blk_sz)
35    draw.line([left, right], fill=(0, 0, 255, 255))
36  for j in xrange(num_col):
37    up = (j * blk_sz, 0)
38    down = (j * blk_sz, height)
39    draw.line([up, down], fill=(0, 0, 255, 255))
40  for i in xrange(num_row):
41    for j in xrange(num_col):
42      center = (j * blk_sz + 0.5 * blk_sz, i * blk_sz + 0.5 * blk_sz)
43      """mf[i,j][0] is the row shift and mf[i,j][1] is the column shift In PIL coordinates, head[0] is x (column shift) and head[1] is y (row shift)."""
44      head = (center[0] + mf[i, j][1], center[1] + mf[i, j][0])
45      draw.line([center, head], fill=(255, 0, 0, 255))
46  return Image.alpha_composite(img_rgba, mf_layer)
47