ref: e2d0e7fe01984a7fb57fee11a44d57a051a2c464
parent: 0ad301e5b04951ecda115a2869475756760503bb
author: Dan Zhu <zxdan@google.com>
date: Thu Aug 22 14:15:07 EDT 2019
Fix some bugs of python code Change-Id: I509cbda24d7d0c8dac75209efa40e24c09a107c5 Exhaust: add exhaust search with neighbor constraint GroundTruth: be able to import motion field variable MotionEST: use new function names Util: be able to set the size of image Change-Id: I36cfdf4b1f28b8190b3ad2be61c241da1347cfc3
--- a/tools/3D-Reconstruction/MotionEST/Exhaust.py
+++ b/tools/3D-Reconstruction/MotionEST/Exhaust.py
@@ -131,3 +131,121 @@
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
--- a/tools/3D-Reconstruction/MotionEST/GroundTruth.py
+++ b/tools/3D-Reconstruction/MotionEST/GroundTruth.py
@@ -1,36 +1,40 @@
-#!/usr/bin/env python
-# coding: utf-8
+#!/ usr / bin / env python
+#coding : utf - 8
import numpy as np
import numpy.linalg as LA
from MotionEST import MotionEST
"""Ground Truth:
- Load in ground truth motion field and mask
+ Load in ground truth motion field and mask
"""
class GroundTruth(MotionEST):
- """
- constructor:
- cur_f: current frame
- ref_f: reference frame
- blk_sz: block size
- gt_path: ground truth motion field file path
+ """constructor:
+
+ cur_f:current
+ frame ref_f:reference
+ frame blk_sz:block size
+ gt_path:ground truth motion field file path
"""
- def __init__(self, cur_f, ref_f, blk_sz, gt_path):
+ def __init__(self, cur_f, ref_f, blk_sz, gt_path, mf=None, mask=None):
self.name = 'ground truth'
super(GroundTruth, self).__init__(cur_f, ref_f, blk_sz)
self.mask = np.zeros((self.num_row, self.num_col), dtype=np.bool)
- with open(gt_path) as gt_file:
- lines = gt_file.readlines()
- for i in xrange(len(lines)):
- info = lines[i].split(';')- for j in xrange(len(info)):
- x, y = info[j].split(',')- #-,- stands for nothing
- if x == '-' or y == '-':
- self.mask[i, -j - 1] = True
- continue
- #the order of original file is flipped on the x axis
- self.mf[i, -j - 1] = np.array([float(y), -float(x)], dtype=np.int)
+ if not gt_path is None:
+ with open(gt_path) as gt_file:
+ lines = gt_file.readlines()
+ for i in xrange(len(lines)):
+ info = lines[i].split(';')+ for j in xrange(len(info)):
+ x, y = info[j].split(',')+ #-, - stands for nothing
+ if x == '-' or y == '-':
+ self.mask[i, -j - 1] = True
+ continue
+ #the order of original file is flipped on the x axis
+ self.mf[i, -j - 1] = np.array([float(y), -float(x)], dtype=np.int)
+ else:
+ self.mf = mf
+ self.mask = mask
--- a/tools/3D-Reconstruction/MotionEST/MotionEST.py
+++ b/tools/3D-Reconstruction/MotionEST/MotionEST.py
@@ -1,5 +1,5 @@
-#!/usr/bin/env python
-# coding: utf-8
+#!/ usr / bin / env python
+#coding : utf - 8
import numpy as np
import numpy.linalg as LA
import matplotlib.pyplot as plt
@@ -20,8 +20,8 @@
self.ref_f = ref_f
self.blk_sz = blk_sz
#convert RGB to YUV
- self.cur_yuv = np.array(self.cur_f.convert('YCbCr'))- self.ref_yuv = np.array(self.ref_f.convert('YCbCr'))+ self.cur_yuv = np.array(self.cur_f.convert('YCbCr'), dtype=np.int)+ self.ref_yuv = np.array(self.ref_f.convert('YCbCr'), dtype=np.int)#frame size
self.width = self.cur_f.size[0]
self.height = self.cur_f.size[1]
@@ -31,10 +31,7 @@
#initialize motion field
self.mf = np.zeros((self.num_row, self.num_col, 2))
- """
- estimation function
- Override by child classes
- """
+ """estimation function Override by child classes"""
def motion_field_estimation(self):
pass
@@ -41,11 +38,11 @@
"""
distortion of a block:
- cur_r: current row
- cur_c: current column
- mv: motion vector
- metric: distortion metric
- """
+ cur_r: current row
+ cur_c: current column
+ mv: motion vector
+ metric: distortion metric
+ """
def block_dist(self, cur_r, cur_c, mv, metric=MSE):
cur_x = cur_c * self.blk_sz
@@ -63,7 +60,7 @@
"""
distortion of motion field
- """
+ """
def distortion(self, mask=None, metric=MSE):
loss = 0
@@ -72,14 +69,11 @@
for j in xrange(self.num_col):
if not mask is None and mask[i, j]:
continue
- loss += self.dist(i, j, self.mf[i, j], metric)
+ loss += self.block_dist(i, j, self.mf[i, j], metric)
count += 1
return loss / count
- """
- evaluation
- compare the difference with ground truth
- """
+ """evaluation compare the difference with ground truth"""
def motion_field_evaluation(self, ground_truth):
loss = 0
@@ -94,11 +88,9 @@
count += 1
return loss / count
- """
- render the motion field
- """
+ """render the motion field"""
- def show(self, ground_truth=None):
+ def show(self, ground_truth=None, size=10):
cur_mf = drawMF(self.cur_f, self.blk_sz, self.mf)
if ground_truth is None:
n_row = 1
@@ -105,7 +97,7 @@
else:
gt_mf = drawMF(self.cur_f, self.blk_sz, ground_truth)
n_row = 2
- plt.figure(figsize=(n_row * 10, 10))
+ plt.figure(figsize=(n_row * size, size * self.height / self.width))
plt.subplot(1, n_row, 1)
plt.imshow(cur_mf)
plt.title('Estimated Motion Field')--- a/tools/3D-Reconstruction/MotionEST/Util.py
+++ b/tools/3D-Reconstruction/MotionEST/Util.py
@@ -32,8 +32,7 @@
for i in xrange(num_row):
for j in xrange(num_col):
center = (j * blk_sz + 0.5 * blk_sz, i * blk_sz + 0.5 * blk_sz)
- """mf[i,j][0] is the row shift and mf[i,j][1] is the column shift In PIL coordinates, head[0] is x (column shift) and head[1] is y (row shift).
- """
+ """mf[i,j][0] is the row shift and mf[i,j][1] is the column shift In PIL coordinates, head[0] is x (column shift) and head[1] is y (row shift)."""
head = (center[0] + mf[i, j][1], center[1] + mf[i, j][0])
draw.line([center, head], fill=(255, 0, 0, 255))
return Image.alpha_composite(img_rgba, mf_layer)