ref: 2157d613c56095de9ec442376b6386ba9fc0e584
dir: /tools/3D-Reconstruction/MotionEST/Exhaust.py/
#!/usr/bin/env python # coding: utf-8 import numpy as np import numpy.linalg as LA from Util import MSE from MotionEST import MotionEST """Exhaust Search:""" class Exhaust(MotionEST): """ Constructor: cur_f: current frame ref_f: reference frame blk_sz: block size wnd_size: search window size metric: metric to compare the blocks distrotion """ def __init__(self, cur_f, ref_f, blk_size, wnd_size, metric=MSE): self.name = 'exhaust' self.wnd_sz = wnd_size self.metric = metric super(Exhaust, self).__init__(cur_f, ref_f, blk_size) """ search method: cur_r: start row cur_c: start column """ def search(self, cur_r, cur_c): min_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) cur_x = cur_c * self.blk_sz cur_y = cur_r * self.blk_sz ref_x = cur_x ref_y = cur_y #search all validate positions and select the one with minimum distortion for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], self.metric) if loss < min_loss: min_loss = loss ref_x = x ref_y = y return ref_x, ref_y def motion_field_estimation(self): for i in xrange(self.num_row): for j in xrange(self.num_col): ref_x, ref_y = self.search(i, j) self.mf[i, j] = np.array( [ref_y - i * self.blk_sz, ref_x - j * self.blk_sz]) """Exhaust with Neighbor Constraint""" class ExhaustNeighbor(MotionEST): """ Constructor: cur_f: current frame ref_f: reference frame blk_sz: block size wnd_size: search window size beta: neigbor loss weight metric: metric to compare the blocks distrotion """ def __init__(self, cur_f, ref_f, blk_size, wnd_size, beta, metric=MSE): self.name = 'exhaust + neighbor' self.wnd_sz = wnd_size self.beta = beta self.metric = metric super(ExhaustNeighbor, self).__init__(cur_f, ref_f, blk_size) self.assign = np.zeros((self.num_row, self.num_col), dtype=np.bool) """ estimate neighbor loss: cur_r: current row cur_c: current column mv: current motion vector """ def neighborLoss(self, cur_r, cur_c, mv): loss = 0 #accumulate difference between current block's motion vector with neighbors' for i, j in {(-1, 0), (1, 0), (0, 1), (0, -1)}: nb_r = cur_r + i nb_c = cur_c + j if 0 <= nb_r < self.num_row and 0 <= nb_c < self.num_col and self.assign[ nb_r, nb_c]: loss += LA.norm(mv - self.mf[nb_r, nb_c]) return loss """ search method: cur_r: start row cur_c: start column """ def search(self, cur_r, cur_c): dist_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) nb_loss = self.neighborLoss(cur_r, cur_c, np.array([0, 0])) min_loss = dist_loss + self.beta * nb_loss cur_x = cur_c * self.blk_sz cur_y = cur_r * self.blk_sz ref_x = cur_x ref_y = cur_y #search all validate positions and select the one with minimum distortion # as well as weighted neighbor loss for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: dist_loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], self.metric) nb_loss = self.neighborLoss(cur_r, cur_c, [y - cur_y, x - cur_x]) loss = dist_loss + self.beta * nb_loss if loss < min_loss: min_loss = loss ref_x = x ref_y = y return ref_x, ref_y def motion_field_estimation(self): for i in xrange(self.num_row): for j in xrange(self.num_col): ref_x, ref_y = self.search(i, j) self.mf[i, j] = np.array( [ref_y - i * self.blk_sz, ref_x - j * self.blk_sz]) self.assign[i, j] = True """Exhaust with Neighbor Constraint and Feature Score""" class ExhaustNeighborFeatureScore(MotionEST): """ Constructor: cur_f: current frame ref_f: reference frame blk_sz: block size wnd_size: search window size beta: neigbor loss weight max_iter: maximum number of iterations metric: metric to compare the blocks distrotion """ def __init__(self, cur_f, ref_f, blk_size, wnd_size, beta=1, max_iter=100, metric=MSE): self.name = 'exhaust + neighbor+feature score' self.wnd_sz = wnd_size self.beta = beta self.metric = metric self.max_iter = max_iter super(ExhaustNeighborFeatureScore, self).__init__(cur_f, ref_f, blk_size) self.fs = self.getFeatureScore() """ get feature score of each block """ def getFeatureScore(self): fs = np.zeros((self.num_row, self.num_col)) for r in xrange(self.num_row): for c in xrange(self.num_col): IxIx = 0 IyIy = 0 IxIy = 0 #get ssd surface for x in xrange(self.blk_sz - 1): for y in xrange(self.blk_sz - 1): ox = c * self.blk_sz + x oy = r * self.blk_sz + y Ix = self.cur_yuv[oy, ox + 1, 0] - self.cur_yuv[oy, ox, 0] Iy = self.cur_yuv[oy + 1, ox, 0] - self.cur_yuv[oy, ox, 0] IxIx += Ix * Ix IyIy += Iy * Iy IxIy += Ix * Iy #get maximum and minimum eigenvalues lambda_max = 0.5 * ((IxIx + IyIy) + np.sqrt(4 * IxIy * IxIy + (IxIx - IyIy)**2)) lambda_min = 0.5 * ((IxIx + IyIy) - np.sqrt(4 * IxIy * IxIy + (IxIx - IyIy)**2)) fs[r, c] = lambda_max * lambda_min / (1e-6 + lambda_max + lambda_min) if fs[r, c] < 0: fs[r, c] = 0 return fs """ do exhaust search """ def search(self, cur_r, cur_c): min_loss = self.block_dist(cur_r, cur_c, [0, 0], self.metric) cur_x = cur_c * self.blk_sz cur_y = cur_r * self.blk_sz ref_x = cur_x ref_y = cur_y #search all validate positions and select the one with minimum distortion for y in xrange(cur_y - self.wnd_sz, cur_y + self.wnd_sz): for x in xrange(cur_x - self.wnd_sz, cur_x + self.wnd_sz): if 0 <= x < self.width - self.blk_sz and 0 <= y < self.height - self.blk_sz: loss = self.block_dist(cur_r, cur_c, [y - cur_y, x - cur_x], self.metric) if loss < min_loss: min_loss = loss ref_x = x ref_y = y return ref_x, ref_y """ add smooth constraint """ def smooth(self, uvs, mvs): sm_uvs = np.zeros(uvs.shape) for r in xrange(self.num_row): for c in xrange(self.num_col): avg_uv = np.array([0.0, 0.0]) for i, j in {(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)}: if 0 <= i < self.num_row and 0 <= j < self.num_col: avg_uv += uvs[i, j] / 6.0 for i, j in {(r - 1, c - 1), (r - 1, c + 1), (r + 1, c - 1), (r + 1, c + 1)}: if 0 <= i < self.num_row and 0 <= j < self.num_col: avg_uv += uvs[i, j] / 12.0 sm_uvs[r, c] = (self.fs[r, c] * mvs[r, c] + self.beta * avg_uv) / ( self.beta + self.fs[r, c]) return sm_uvs def motion_field_estimation(self): #get matching results mvs = np.zeros(self.mf.shape) for r in xrange(self.num_row): for c in xrange(self.num_col): ref_x, ref_y = self.search(r, c) mvs[r, c] = np.array([ref_y - r * self.blk_sz, ref_x - c * self.blk_sz]) #add smoothness constraint uvs = np.zeros(self.mf.shape) for _ in xrange(self.max_iter): uvs = self.smooth(uvs, mvs) self.mf = uvs