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