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 Util import MSE 14from MotionEST import MotionEST 15"""Search & Smooth Model with Adapt Weights""" 16 17 18class SearchSmoothAdapt(MotionEST): 19 """ 20 Constructor: 21 cur_f: current frame 22 ref_f: reference frame 23 blk_sz: block size 24 wnd_size: search window size 25 beta: neigbor loss weight 26 max_iter: maximum number of iterations 27 metric: metric to compare the blocks distrotion 28 """ 29 30 def __init__(self, cur_f, ref_f, blk_size, search, max_iter=100): 31 self.search = search 32 self.max_iter = max_iter 33 super(SearchSmoothAdapt, self).__init__(cur_f, ref_f, blk_size) 34 35 """ 36 get local diffiencial of refernce 37 """ 38 39 def getRefLocalDiff(self, mvs): 40 m, n = self.num_row, self.num_col 41 localDiff = [[] for _ in xrange(m)] 42 blk_sz = self.blk_sz 43 for r in xrange(m): 44 for c in xrange(n): 45 I_row = 0 46 I_col = 0 47 #get ssd surface 48 count = 0 49 center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz, 50 c * blk_sz:(c + 1) * blk_sz, 0] 51 ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz) 52 tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz) 53 target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0] 54 for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}: 55 if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz: 56 nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0] 57 I_row += np.sum(np.abs(nb - center)) - np.sum( 58 np.abs(target - center)) 59 count += 1 60 I_row //= (count * blk_sz * blk_sz) 61 count = 0 62 for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}: 63 if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz: 64 nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0] 65 I_col += np.sum(np.abs(nb - center)) - np.sum( 66 np.abs(target - center)) 67 count += 1 68 I_col //= (count * blk_sz * blk_sz) 69 localDiff[r].append( 70 np.array([[I_row * I_row, I_row * I_col], 71 [I_col * I_row, I_col * I_col]])) 72 return localDiff 73 74 """ 75 add smooth constraint 76 """ 77 78 def smooth(self, uvs, mvs): 79 sm_uvs = np.zeros(uvs.shape) 80 blk_sz = self.blk_sz 81 for r in xrange(self.num_row): 82 for c in xrange(self.num_col): 83 nb_uv = np.array([0.0, 0.0]) 84 for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}: 85 if 0 <= i < self.num_row and 0 <= j < self.num_col: 86 nb_uv += uvs[i, j] / 6.0 87 else: 88 nb_uv += uvs[r, c] / 6.0 89 for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1), 90 (r + 1, c + 1)}: 91 if 0 <= i < self.num_row and 0 <= j < self.num_col: 92 nb_uv += uvs[i, j] / 12.0 93 else: 94 nb_uv += uvs[r, c] / 12.0 95 ssd_nb = self.block_dist(r, c, self.blk_sz * nb_uv) 96 mv = mvs[r, c] 97 ssd_mv = self.block_dist(r, c, mv) 98 alpha = (ssd_nb - ssd_mv) / (ssd_mv + 1e-6) 99 M = alpha * self.localDiff[r][c] 100 P = M + np.identity(2) 101 inv_P = LA.inv(P) 102 sm_uvs[r, c] = np.dot(inv_P, nb_uv) + np.dot( 103 np.matmul(inv_P, M), mv / blk_sz) 104 return sm_uvs 105 106 def block_matching(self): 107 self.search.motion_field_estimation() 108 109 def motion_field_estimation(self): 110 self.localDiff = self.getRefLocalDiff(self.search.mf) 111 #get matching results 112 mvs = self.search.mf 113 #add smoothness constraint 114 uvs = mvs / self.blk_sz 115 for _ in xrange(self.max_iter): 116 uvs = self.smooth(uvs, mvs) 117 self.mf = uvs * self.blk_sz 118 119 120"""Search & Smooth Model with Fixed Weights""" 121 122 123class SearchSmoothFix(MotionEST): 124 """ 125 Constructor: 126 cur_f: current frame 127 ref_f: reference frame 128 blk_sz: block size 129 wnd_size: search window size 130 beta: neigbor loss weight 131 max_iter: maximum number of iterations 132 metric: metric to compare the blocks distrotion 133 """ 134 135 def __init__(self, cur_f, ref_f, blk_size, search, beta, max_iter=100): 136 self.search = search 137 self.max_iter = max_iter 138 self.beta = beta 139 super(SearchSmoothFix, self).__init__(cur_f, ref_f, blk_size) 140 141 """ 142 get local diffiencial of refernce 143 """ 144 145 def getRefLocalDiff(self, mvs): 146 m, n = self.num_row, self.num_col 147 localDiff = [[] for _ in xrange(m)] 148 blk_sz = self.blk_sz 149 for r in xrange(m): 150 for c in xrange(n): 151 I_row = 0 152 I_col = 0 153 #get ssd surface 154 count = 0 155 center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz, 156 c * blk_sz:(c + 1) * blk_sz, 0] 157 ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz) 158 tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz) 159 target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0] 160 for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}: 161 if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz: 162 nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0] 163 I_row += np.sum(np.abs(nb - center)) - np.sum( 164 np.abs(target - center)) 165 count += 1 166 I_row //= (count * blk_sz * blk_sz) 167 count = 0 168 for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}: 169 if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz: 170 nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0] 171 I_col += np.sum(np.abs(nb - center)) - np.sum( 172 np.abs(target - center)) 173 count += 1 174 I_col //= (count * blk_sz * blk_sz) 175 localDiff[r].append( 176 np.array([[I_row * I_row, I_row * I_col], 177 [I_col * I_row, I_col * I_col]])) 178 return localDiff 179 180 """ 181 add smooth constraint 182 """ 183 184 def smooth(self, uvs, mvs): 185 sm_uvs = np.zeros(uvs.shape) 186 blk_sz = self.blk_sz 187 for r in xrange(self.num_row): 188 for c in xrange(self.num_col): 189 nb_uv = np.array([0.0, 0.0]) 190 for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}: 191 if 0 <= i < self.num_row and 0 <= j < self.num_col: 192 nb_uv += uvs[i, j] / 6.0 193 else: 194 nb_uv += uvs[r, c] / 6.0 195 for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1), 196 (r + 1, c + 1)}: 197 if 0 <= i < self.num_row and 0 <= j < self.num_col: 198 nb_uv += uvs[i, j] / 12.0 199 else: 200 nb_uv += uvs[r, c] / 12.0 201 mv = mvs[r, c] / blk_sz 202 M = self.localDiff[r][c] 203 P = M + self.beta * np.identity(2) 204 inv_P = LA.inv(P) 205 sm_uvs[r, c] = np.dot(inv_P, self.beta * nb_uv) + np.dot( 206 np.matmul(inv_P, M), mv) 207 return sm_uvs 208 209 def block_matching(self): 210 self.search.motion_field_estimation() 211 212 def motion_field_estimation(self): 213 #get local structure 214 self.localDiff = self.getRefLocalDiff(self.search.mf) 215 #get matching results 216 mvs = self.search.mf 217 #add smoothness constraint 218 uvs = mvs / self.blk_sz 219 for _ in xrange(self.max_iter): 220 uvs = self.smooth(uvs, mvs) 221 self.mf = uvs * self.blk_sz 222