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"""Horn & Schunck Model"""
18
19
20class HornSchunck(MotionEST):
21  """
22    constructor:
23        cur_f: current frame
24        ref_f: reference frame
25        blk_sz: block size
26        alpha: smooth constrain weight
27        sigma: gaussian blur parameter
28    """
29
30  def __init__(self, cur_f, ref_f, blk_sz, alpha, sigma, max_iter=100):
31    super(HornSchunck, self).__init__(cur_f, ref_f, blk_sz)
32    self.cur_I, self.ref_I = self.getIntensity()
33    #perform gaussian blur to smooth the intensity
34    self.cur_I = gaussian_filter(self.cur_I, sigma=sigma)
35    self.ref_I = gaussian_filter(self.ref_I, sigma=sigma)
36    self.alpha = alpha
37    self.max_iter = max_iter
38    self.Ix, self.Iy, self.It = self.intensityDiff()
39
40  """
41    Build Frame Intensity
42    """
43
44  def getIntensity(self):
45    cur_I = np.zeros((self.num_row, self.num_col))
46    ref_I = np.zeros((self.num_row, self.num_col))
47    #use average intensity as block's intensity
48    for i in xrange(self.num_row):
49      for j in xrange(self.num_col):
50        r = i * self.blk_sz
51        c = j * self.blk_sz
52        cur_I[i, j] = np.mean(self.cur_yuv[r:r + self.blk_sz, c:c + self.blk_sz,
53                                           0])
54        ref_I[i, j] = np.mean(self.ref_yuv[r:r + self.blk_sz, c:c + self.blk_sz,
55                                           0])
56    return cur_I, ref_I
57
58  """
59    Get First Order Derivative
60    """
61
62  def intensityDiff(self):
63    Ix = np.zeros((self.num_row, self.num_col))
64    Iy = np.zeros((self.num_row, self.num_col))
65    It = np.zeros((self.num_row, self.num_col))
66    sz = self.blk_sz
67    for i in xrange(self.num_row - 1):
68      for j in xrange(self.num_col - 1):
69        """
70                Ix:
71                (i  ,j) <--- (i  ,j+1)
72                (i+1,j) <--- (i+1,j+1)
73                """
74        count = 0
75        for r, c in {(i, j + 1), (i + 1, j + 1)}:
76          if 0 <= r < self.num_row and 0 < c < self.num_col:
77            Ix[i, j] += (
78                self.cur_I[r, c] - self.cur_I[r, c - 1] + self.ref_I[r, c] -
79                self.ref_I[r, c - 1])
80            count += 2
81        Ix[i, j] /= count
82        """
83                Iy:
84                (i  ,j)      (i  ,j+1)
85                   ^             ^
86                   |             |
87                (i+1,j)      (i+1,j+1)
88                """
89        count = 0
90        for r, c in {(i + 1, j), (i + 1, j + 1)}:
91          if 0 < r < self.num_row and 0 <= c < self.num_col:
92            Iy[i, j] += (
93                self.cur_I[r, c] - self.cur_I[r - 1, c] + self.ref_I[r, c] -
94                self.ref_I[r - 1, c])
95            count += 2
96        Iy[i, j] /= count
97        count = 0
98        #It:
99        for r in xrange(i, i + 2):
100          for c in xrange(j, j + 2):
101            if 0 <= r < self.num_row and 0 <= c < self.num_col:
102              It[i, j] += (self.ref_I[r, c] - self.cur_I[r, c])
103              count += 1
104        It[i, j] /= count
105    return Ix, Iy, It
106
107  """
108    Get weighted average of neighbor motion vectors
109    for evaluation of laplacian
110    """
111
112  def averageMV(self):
113    avg = np.zeros((self.num_row, self.num_col, 2))
114    """
115        1/12 ---  1/6 --- 1/12
116         |         |       |
117        1/6  --- -1/8 --- 1/6
118         |         |       |
119        1/12 ---  1/6 --- 1/12
120        """
121    for i in xrange(self.num_row):
122      for j in xrange(self.num_col):
123        for r, c in {(-1, 0), (1, 0), (0, -1), (0, 1)}:
124          if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
125            avg[i, j] += self.mf[i + r, j + c] / 6.0
126        for r, c in {(-1, -1), (-1, 1), (1, -1), (1, 1)}:
127          if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
128            avg[i, j] += self.mf[i + r, j + c] / 12.0
129    return avg
130
131  def motion_field_estimation(self):
132    count = 0
133    """
134        u_{n+1} = ~u_n - Ix(Ix.~u_n+Iy.~v+It)/(IxIx+IyIy+alpha^2)
135        v_{n+1} = ~v_n - Iy(Ix.~u_n+Iy.~v+It)/(IxIx+IyIy+alpha^2)
136        """
137    denom = self.alpha**2 + np.power(self.Ix, 2) + np.power(self.Iy, 2)
138    while count < self.max_iter:
139      avg = self.averageMV()
140      self.mf[:, :, 1] = avg[:, :, 1] - self.Ix * (
141          self.Ix * avg[:, :, 1] + self.Iy * avg[:, :, 0] + self.It) / denom
142      self.mf[:, :, 0] = avg[:, :, 0] - self.Iy * (
143          self.Ix * avg[:, :, 1] + self.Iy * avg[:, :, 0] + self.It) / denom
144      count += 1
145    self.mf *= self.blk_sz
146
147  def motion_field_estimation_mat(self):
148    row_idx = []
149    col_idx = []
150    data = []
151
152    N = 2 * self.num_row * self.num_col
153    b = np.zeros((N, 1))
154    for i in xrange(self.num_row):
155      for j in xrange(self.num_col):
156        """(IxIx+alpha^2)u+IxIy.v-alpha^2~u IxIy.u+(IyIy+alpha^2)v-alpha^2~v"""
157        u_idx = i * 2 * self.num_col + 2 * j
158        v_idx = u_idx + 1
159        b[u_idx, 0] = -self.Ix[i, j] * self.It[i, j]
160        b[v_idx, 0] = -self.Iy[i, j] * self.It[i, j]
161        #u: (IxIx+alpha^2)u
162        row_idx.append(u_idx)
163        col_idx.append(u_idx)
164        data.append(self.Ix[i, j] * self.Ix[i, j] + self.alpha**2)
165        #IxIy.v
166        row_idx.append(u_idx)
167        col_idx.append(v_idx)
168        data.append(self.Ix[i, j] * self.Iy[i, j])
169
170        #v: IxIy.u
171        row_idx.append(v_idx)
172        col_idx.append(u_idx)
173        data.append(self.Ix[i, j] * self.Iy[i, j])
174        #(IyIy+alpha^2)v
175        row_idx.append(v_idx)
176        col_idx.append(v_idx)
177        data.append(self.Iy[i, j] * self.Iy[i, j] + self.alpha**2)
178
179        #-alpha^2~u
180        #-alpha^2~v
181        for r, c in {(-1, 0), (1, 0), (0, -1), (0, 1)}:
182          if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
183            u_nb = (i + r) * 2 * self.num_col + 2 * (j + c)
184            v_nb = u_nb + 1
185
186            row_idx.append(u_idx)
187            col_idx.append(u_nb)
188            data.append(-1 * self.alpha**2 / 6.0)
189
190            row_idx.append(v_idx)
191            col_idx.append(v_nb)
192            data.append(-1 * self.alpha**2 / 6.0)
193        for r, c in {(-1, -1), (-1, 1), (1, -1), (1, 1)}:
194          if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
195            u_nb = (i + r) * 2 * self.num_col + 2 * (j + c)
196            v_nb = u_nb + 1
197
198            row_idx.append(u_idx)
199            col_idx.append(u_nb)
200            data.append(-1 * self.alpha**2 / 12.0)
201
202            row_idx.append(v_idx)
203            col_idx.append(v_nb)
204            data.append(-1 * self.alpha**2 / 12.0)
205    M = csc_matrix((data, (row_idx, col_idx)), shape=(N, N))
206    M_inv = inv(M)
207    uv = M_inv.dot(b)
208
209    for i in xrange(self.num_row):
210      for j in xrange(self.num_col):
211        self.mf[i, j, 0] = uv[i * 2 * self.num_col + 2 * j + 1, 0] * self.blk_sz
212        self.mf[i, j, 1] = uv[i * 2 * self.num_col + 2 * j, 0] * self.blk_sz
213