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"""Exhaust Search:""" 16 17 18class Exhaust(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 metric: metric to compare the blocks distrotion 26 """ 27 28 def __init__(self, cur_f, ref_f, blk_size, wnd_size, metric=MSE): 29 self.name = 'exhaust' 30 self.wnd_sz = wnd_size 31 self.metric = metric 32 super(Exhaust, self).__init__(cur_f, ref_f, blk_size) 33 34 """ 35 search method: 36 cur_r: start row 37 cur_c: start column 38 """ 39 40 def search(self, cur_r, cur_c): 41 min_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) 42 cur_x = cur_c * self.blk_sz 43 cur_y = cur_r * self.blk_sz 44 ref_x = cur_x 45 ref_y = cur_y 46 #search all validate positions and select the one with minimum distortion 47 for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): 48 for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): 49 if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: 50 loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], 51 self.metric) 52 if loss < min_loss: 53 min_loss = loss 54 ref_x = x 55 ref_y = y 56 return ref_x, ref_y 57 58 def motion_field_estimation(self): 59 for i in xrange(self.num_row): 60 for j in xrange(self.num_col): 61 ref_x, ref_y = self.search(i, j) 62 self.mf[i, j] = np.array( 63 [ref_y - i * self.blk_sz, ref_x - j * self.blk_sz]) 64 65 66"""Exhaust with Neighbor Constraint""" 67 68 69class ExhaustNeighbor(MotionEST): 70 """ 71 Constructor: 72 cur_f: current frame 73 ref_f: reference frame 74 blk_sz: block size 75 wnd_size: search window size 76 beta: neigbor loss weight 77 metric: metric to compare the blocks distrotion 78 """ 79 80 def __init__(self, cur_f, ref_f, blk_size, wnd_size, beta, metric=MSE): 81 self.name = 'exhaust + neighbor' 82 self.wnd_sz = wnd_size 83 self.beta = beta 84 self.metric = metric 85 super(ExhaustNeighbor, self).__init__(cur_f, ref_f, blk_size) 86 self.assign = np.zeros((self.num_row, self.num_col), dtype=np.bool) 87 88 """ 89 estimate neighbor loss: 90 cur_r: current row 91 cur_c: current column 92 mv: current motion vector 93 """ 94 95 def neighborLoss(self, cur_r, cur_c, mv): 96 loss = 0 97 #accumulate difference between current block's motion vector with neighbors' 98 for i, j in {(-1, 0), (1, 0), (0, 1), (0, -1)}: 99 nb_r = cur_r + i 100 nb_c = cur_c + j 101 if 0 <= nb_r < self.num_row and 0 <= nb_c < self.num_col and self.assign[ 102 nb_r, nb_c]: 103 loss += LA.norm(mv - self.mf[nb_r, nb_c]) 104 return loss 105 106 """ 107 search method: 108 cur_r: start row 109 cur_c: start column 110 """ 111 112 def search(self, cur_r, cur_c): 113 dist_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) 114 nb_loss = self.neighborLoss(cur_r, cur_c, np.array([0, 0])) 115 min_loss = dist_loss + self.beta * nb_loss 116 cur_x = cur_c * self.blk_sz 117 cur_y = cur_r * self.blk_sz 118 ref_x = cur_x 119 ref_y = cur_y 120 #search all validate positions and select the one with minimum distortion 121 # as well as weighted neighbor loss 122 for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): 123 for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): 124 if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: 125 dist_loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], 126 self.metric) 127 nb_loss = self.neighborLoss(cur_r, cur_c, [y - cur_y, x - cur_x]) 128 loss = dist_loss + self.beta * nb_loss 129 if loss < min_loss: 130 min_loss = loss 131 ref_x = x 132 ref_y = y 133 return ref_x, ref_y 134 135 def motion_field_estimation(self): 136 for i in xrange(self.num_row): 137 for j in xrange(self.num_col): 138 ref_x, ref_y = self.search(i, j) 139 self.mf[i, j] = np.array( 140 [ref_y - i * self.blk_sz, ref_x - j * self.blk_sz]) 141 self.assign[i, j] = True 142 143 144"""Exhaust with Neighbor Constraint and Feature Score""" 145 146 147class ExhaustNeighborFeatureScore(MotionEST): 148 """ 149 Constructor: 150 cur_f: current frame 151 ref_f: reference frame 152 blk_sz: block size 153 wnd_size: search window size 154 beta: neigbor loss weight 155 max_iter: maximum number of iterations 156 metric: metric to compare the blocks distrotion 157 """ 158 159 def __init__(self, 160 cur_f, 161 ref_f, 162 blk_size, 163 wnd_size, 164 beta=1, 165 max_iter=100, 166 metric=MSE): 167 self.name = 'exhaust + neighbor+feature score' 168 self.wnd_sz = wnd_size 169 self.beta = beta 170 self.metric = metric 171 self.max_iter = max_iter 172 super(ExhaustNeighborFeatureScore, self).__init__(cur_f, ref_f, blk_size) 173 self.fs = self.getFeatureScore() 174 175 """ 176 get feature score of each block 177 """ 178 179 def getFeatureScore(self): 180 fs = np.zeros((self.num_row, self.num_col)) 181 for r in xrange(self.num_row): 182 for c in xrange(self.num_col): 183 IxIx = 0 184 IyIy = 0 185 IxIy = 0 186 #get ssd surface 187 for x in xrange(self.blk_sz - 1): 188 for y in xrange(self.blk_sz - 1): 189 ox = c * self.blk_sz + x 190 oy = r * self.blk_sz + y 191 Ix = self.cur_yuv[oy, ox + 1, 0] - self.cur_yuv[oy, ox, 0] 192 Iy = self.cur_yuv[oy + 1, ox, 0] - self.cur_yuv[oy, ox, 0] 193 IxIx += Ix * Ix 194 IyIy += Iy * Iy 195 IxIy += Ix * Iy 196 #get maximum and minimum eigenvalues 197 lambda_max = 0.5 * ((IxIx + IyIy) + np.sqrt(4 * IxIy * IxIy + 198 (IxIx - IyIy)**2)) 199 lambda_min = 0.5 * ((IxIx + IyIy) - np.sqrt(4 * IxIy * IxIy + 200 (IxIx - IyIy)**2)) 201 fs[r, c] = lambda_max * lambda_min / (1e-6 + lambda_max + lambda_min) 202 if fs[r, c] < 0: 203 fs[r, c] = 0 204 return fs 205 206 """ 207 do exhaust search 208 """ 209 210 def search(self, cur_r, cur_c): 211 min_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) 212 cur_x = cur_c * self.blk_sz 213 cur_y = cur_r * self.blk_sz 214 ref_x = cur_x 215 ref_y = cur_y 216 #search all validate positions and select the one with minimum distortion 217 for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): 218 for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): 219 if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: 220 loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], 221 self.metric) 222 if loss < min_loss: 223 min_loss = loss 224 ref_x = x 225 ref_y = y 226 return ref_x, ref_y 227 228 """ 229 add smooth constraint 230 """ 231 232 def smooth(self, uvs, mvs): 233 sm_uvs = np.zeros(uvs.shape) 234 for r in xrange(self.num_row): 235 for c in xrange(self.num_col): 236 avg_uv = np.array([0.0, 0.0]) 237 for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}: 238 if 0 <= i < self.num_row and 0 <= j < self.num_col: 239 avg_uv += uvs[i, j] / 6.0 240 for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1), 241 (r + 1, c + 1)}: 242 if 0 <= i < self.num_row and 0 <= j < self.num_col: 243 avg_uv += uvs[i, j] / 12.0 244 sm_uvs[r, c] = (self.fs[r, c] * mvs[r, c] + self.beta * avg_uv) / ( 245 self.beta + self.fs[r, c]) 246 return sm_uvs 247 248 def motion_field_estimation(self): 249 #get matching results 250 mvs = np.zeros(self.mf.shape) 251 for r in xrange(self.num_row): 252 for c in xrange(self.num_col): 253 ref_x, ref_y = self.search(r, c) 254 mvs[r, c] = np.array([ref_y - r * self.blk_sz, ref_x - c * self.blk_sz]) 255 #add smoothness constraint 256 uvs = np.zeros(self.mf.shape) 257 for _ in xrange(self.max_iter): 258 uvs = self.smooth(uvs, mvs) 259 self.mf = uvs 260