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