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