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

4 files changed, 161 insertions, 48 deletions

diff --git a/tools/3D-Reconstruction/MotionEST/Exhaust.py b/tools/3D-Reconstruction/MotionEST/Exhaust.py
index 83ca157d4..97cfc41de 100644
--- a/tools/3D-Reconstruction/MotionEST/Exhaust.py
+++ b/tools/3D-Reconstruction/MotionEST/Exhaust.py
@@ -131,3 +131,121 @@ class ExhaustNeighbor(MotionEST):
         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
diff --git a/tools/3D-Reconstruction/MotionEST/GroundTruth.py b/tools/3D-Reconstruction/MotionEST/GroundTruth.py
index 61b4bef42..94bd87a61 100644
--- 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
diff --git a/tools/3D-Reconstruction/MotionEST/MotionEST.py b/tools/3D-Reconstruction/MotionEST/MotionEST.py
index 68cf7e743..563fe1927 100644
--- 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 @@ class MotionEST(object):
     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,21 +31,18 @@ class MotionEST(object):
     #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
 
   """
     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 @@ class MotionEST(object):
 
   """
     distortion of motion field
-    """
+  """
 
   def distortion(self, mask=None, metric=MSE):
     loss = 0
@@ -72,14 +69,11 @@ class MotionEST(object):
       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,18 +88,16 @@ class MotionEST(object):
         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
     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')
diff --git a/tools/3D-Reconstruction/MotionEST/Util.py b/tools/3D-Reconstruction/MotionEST/Util.py
index d52e8a5c4..a5265f44e 100644
--- a/tools/3D-Reconstruction/MotionEST/Util.py
+++ b/tools/3D-Reconstruction/MotionEST/Util.py
@@ -32,8 +32,7 @@ def drawMF(img, blk_sz, mf):
   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)