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
13from scipy.ndimage.filters import gaussian_filter
14from scipy.sparse import csc_matrix
15from scipy.sparse.linalg import inv
16from MotionEST import MotionEST
17"""Anandan Model"""
18
19
20class Anandan(MotionEST):
21  """
22    constructor:
23        cur_f: current frame
24        ref_f: reference frame
25        blk_sz: block size
26        beta: smooth constrain weight
27        k1,k2,k3: confidence coefficients
28        max_iter: maximum number of iterations
29    """
30
31  def __init__(self, cur_f, ref_f, blk_sz, beta, k1, k2, k3, max_iter=100):
32    super(Anandan, self).__init__(cur_f, ref_f, blk_sz)
33    self.levels = int(np.log2(blk_sz))
34    self.intensity_hierarchy()
35    self.c_maxs = []
36    self.c_mins = []
37    self.e_maxs = []
38    self.e_mins = []
39    for l in xrange(self.levels + 1):
40      c_max, c_min, e_max, e_min = self.get_curvature(self.cur_Is[l])
41      self.c_maxs.append(c_max)
42      self.c_mins.append(c_min)
43      self.e_maxs.append(e_max)
44      self.e_mins.append(e_min)
45    self.beta = beta
46    self.k1, self.k2, self.k3 = k1, k2, k3
47    self.max_iter = max_iter
48
49  """
50    build intensity hierarchy
51    """
52
53  def intensity_hierarchy(self):
54    level = 0
55    self.cur_Is = []
56    self.ref_Is = []
57    #build each level itensity by using gaussian filters
58    while level <= self.levels:
59      cur_I = gaussian_filter(self.cur_yuv[:, :, 0], sigma=(2**level) * 0.56)
60      ref_I = gaussian_filter(self.ref_yuv[:, :, 0], sigma=(2**level) * 0.56)
61      self.ref_Is.append(ref_I)
62      self.cur_Is.append(cur_I)
63      level += 1
64
65  """
66    get curvature of each block
67    """
68
69  def get_curvature(self, I):
70    c_max = np.zeros((self.num_row, self.num_col))
71    c_min = np.zeros((self.num_row, self.num_col))
72    e_max = np.zeros((self.num_row, self.num_col, 2))
73    e_min = np.zeros((self.num_row, self.num_col, 2))
74    for r in xrange(self.num_row):
75      for c in xrange(self.num_col):
76        h11, h12, h21, h22 = 0, 0, 0, 0
77        for i in xrange(r * self.blk_sz, r * self.blk_sz + self.blk_sz):
78          for j in xrange(c * self.blk_sz, c * self.blk_sz + self.blk_sz):
79            if 0 <= i < self.height - 1 and 0 <= j < self.width - 1:
80              Ix = I[i][j + 1] - I[i][j]
81              Iy = I[i + 1][j] - I[i][j]
82              h11 += Iy * Iy
83              h12 += Ix * Iy
84              h21 += Ix * Iy
85              h22 += Ix * Ix
86        U, S, _ = LA.svd(np.array([[h11, h12], [h21, h22]]))
87        c_max[r, c], c_min[r, c] = S[0], S[1]
88        e_max[r, c] = U[:, 0]
89        e_min[r, c] = U[:, 1]
90    return c_max, c_min, e_max, e_min
91
92  """
93    get ssd of motion vector:
94      cur_I: current intensity
95      ref_I: reference intensity
96      center: current position
97      mv: motion vector
98    """
99
100  def get_ssd(self, cur_I, ref_I, center, mv):
101    ssd = 0
102    for r in xrange(int(center[0]), int(center[0]) + self.blk_sz):
103      for c in xrange(int(center[1]), int(center[1]) + self.blk_sz):
104        if 0 <= r < self.height and 0 <= c < self.width:
105          tr, tc = r + int(mv[0]), c + int(mv[1])
106          if 0 <= tr < self.height and 0 <= tc < self.width:
107            ssd += (ref_I[tr, tc] - cur_I[r, c])**2
108          else:
109            ssd += cur_I[r, c]**2
110    return ssd
111
112  """
113    get region match of level l
114      l: current level
115      last_mvs: matchine results of last level
116      radius: movenment radius
117    """
118
119  def region_match(self, l, last_mvs, radius):
120    mvs = np.zeros((self.num_row, self.num_col, 2))
121    min_ssds = np.zeros((self.num_row, self.num_col))
122    for r in xrange(self.num_row):
123      for c in xrange(self.num_col):
124        center = np.array([r * self.blk_sz, c * self.blk_sz])
125        #use overlap hierarchy policy
126        init_mvs = []
127        if last_mvs is None:
128          init_mvs = [np.array([0, 0])]
129        else:
130          for i, j in {(r, c), (r, c + 1), (r + 1, c), (r + 1, c + 1)}:
131            if 0 <= i < last_mvs.shape[0] and 0 <= j < last_mvs.shape[1]:
132              init_mvs.append(last_mvs[i, j])
133        #use last matching results as the start position as current level
134        min_ssd = None
135        min_mv = None
136        for init_mv in init_mvs:
137          for i in xrange(-2, 3):
138            for j in xrange(-2, 3):
139              mv = init_mv + np.array([i, j]) * radius
140              ssd = self.get_ssd(self.cur_Is[l], self.ref_Is[l], center, mv)
141              if min_ssd is None or ssd < min_ssd:
142                min_ssd = ssd
143                min_mv = mv
144        min_ssds[r, c] = min_ssd
145        mvs[r, c] = min_mv
146    return mvs, min_ssds
147
148  """
149    smooth motion field based on neighbor constraint
150      uvs: current estimation
151      mvs: matching results
152      min_ssds: minimum ssd of matching results
153      l: current level
154    """
155
156  def smooth(self, uvs, mvs, min_ssds, l):
157    sm_uvs = np.zeros((self.num_row, self.num_col, 2))
158    c_max = self.c_maxs[l]
159    c_min = self.c_mins[l]
160    e_max = self.e_maxs[l]
161    e_min = self.e_mins[l]
162    for r in xrange(self.num_row):
163      for c in xrange(self.num_col):
164        w_max = c_max[r, c] / (
165            self.k1 + self.k2 * min_ssds[r, c] + self.k3 * c_max[r, c])
166        w_min = c_min[r, c] / (
167            self.k1 + self.k2 * min_ssds[r, c] + self.k3 * c_min[r, c])
168        w = w_max * w_min / (w_max + w_min + 1e-6)
169        if w < 0:
170          w = 0
171        avg_uv = np.array([0.0, 0.0])
172        for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}:
173          if 0 <= i < self.num_row and 0 <= j < self.num_col:
174            avg_uv += 0.25 * uvs[i, j]
175        sm_uvs[r, c] = (w * w * mvs[r, c] + self.beta * avg_uv) / (
176            self.beta + w * w)
177    return sm_uvs
178
179  """
180    motion field estimation
181    """
182
183  def motion_field_estimation(self):
184    last_mvs = None
185    for l in xrange(self.levels, -1, -1):
186      mvs, min_ssds = self.region_match(l, last_mvs, 2**l)
187      uvs = np.zeros(mvs.shape)
188      for _ in xrange(self.max_iter):
189        uvs = self.smooth(uvs, mvs, min_ssds, l)
190      last_mvs = uvs
191    for r in xrange(self.num_row):
192      for c in xrange(self.num_col):
193        self.mf[r, c] = uvs[r, c]
194