222 lines
7.5 KiB
Python
222 lines
7.5 KiB
Python
## Copyright (c) 2020 The WebM project authors. All Rights Reserved.
|
|
##
|
|
## Use of this source code is governed by a BSD-style license
|
|
## that can be found in the LICENSE file in the root of the source
|
|
## tree. An additional intellectual property rights grant can be found
|
|
## in the file PATENTS. All contributing project authors may
|
|
## be found in the AUTHORS file in the root of the source tree.
|
|
##
|
|
|
|
# coding: utf-8
|
|
import numpy as np
|
|
import numpy.linalg as LA
|
|
from Util import MSE
|
|
from MotionEST import MotionEST
|
|
"""Search & Smooth Model with Adapt Weights"""
|
|
|
|
|
|
class SearchSmoothAdapt(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, search, max_iter=100):
|
|
self.search = search
|
|
self.max_iter = max_iter
|
|
super(SearchSmoothAdapt, self).__init__(cur_f, ref_f, blk_size)
|
|
|
|
"""
|
|
get local diffiencial of refernce
|
|
"""
|
|
|
|
def getRefLocalDiff(self, mvs):
|
|
m, n = self.num_row, self.num_col
|
|
localDiff = [[] for _ in xrange(m)]
|
|
blk_sz = self.blk_sz
|
|
for r in xrange(m):
|
|
for c in xrange(n):
|
|
I_row = 0
|
|
I_col = 0
|
|
#get ssd surface
|
|
count = 0
|
|
center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz,
|
|
c * blk_sz:(c + 1) * blk_sz, 0]
|
|
ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz)
|
|
tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz)
|
|
target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0]
|
|
for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}:
|
|
if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
|
|
nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
|
|
I_row += np.sum(np.abs(nb - center)) - np.sum(
|
|
np.abs(target - center))
|
|
count += 1
|
|
I_row //= (count * blk_sz * blk_sz)
|
|
count = 0
|
|
for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}:
|
|
if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
|
|
nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
|
|
I_col += np.sum(np.abs(nb - center)) - np.sum(
|
|
np.abs(target - center))
|
|
count += 1
|
|
I_col //= (count * blk_sz * blk_sz)
|
|
localDiff[r].append(
|
|
np.array([[I_row * I_row, I_row * I_col],
|
|
[I_col * I_row, I_col * I_col]]))
|
|
return localDiff
|
|
|
|
"""
|
|
add smooth constraint
|
|
"""
|
|
|
|
def smooth(self, uvs, mvs):
|
|
sm_uvs = np.zeros(uvs.shape)
|
|
blk_sz = self.blk_sz
|
|
for r in xrange(self.num_row):
|
|
for c in xrange(self.num_col):
|
|
nb_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:
|
|
nb_uv += uvs[i, j] / 6.0
|
|
else:
|
|
nb_uv += uvs[r, c] / 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:
|
|
nb_uv += uvs[i, j] / 12.0
|
|
else:
|
|
nb_uv += uvs[r, c] / 12.0
|
|
ssd_nb = self.block_dist(r, c, self.blk_sz * nb_uv)
|
|
mv = mvs[r, c]
|
|
ssd_mv = self.block_dist(r, c, mv)
|
|
alpha = (ssd_nb - ssd_mv) / (ssd_mv + 1e-6)
|
|
M = alpha * self.localDiff[r][c]
|
|
P = M + np.identity(2)
|
|
inv_P = LA.inv(P)
|
|
sm_uvs[r, c] = np.dot(inv_P, nb_uv) + np.dot(
|
|
np.matmul(inv_P, M), mv / blk_sz)
|
|
return sm_uvs
|
|
|
|
def block_matching(self):
|
|
self.search.motion_field_estimation()
|
|
|
|
def motion_field_estimation(self):
|
|
self.localDiff = self.getRefLocalDiff(self.search.mf)
|
|
#get matching results
|
|
mvs = self.search.mf
|
|
#add smoothness constraint
|
|
uvs = mvs / self.blk_sz
|
|
for _ in xrange(self.max_iter):
|
|
uvs = self.smooth(uvs, mvs)
|
|
self.mf = uvs * self.blk_sz
|
|
|
|
|
|
"""Search & Smooth Model with Fixed Weights"""
|
|
|
|
|
|
class SearchSmoothFix(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, search, beta, max_iter=100):
|
|
self.search = search
|
|
self.max_iter = max_iter
|
|
self.beta = beta
|
|
super(SearchSmoothFix, self).__init__(cur_f, ref_f, blk_size)
|
|
|
|
"""
|
|
get local diffiencial of refernce
|
|
"""
|
|
|
|
def getRefLocalDiff(self, mvs):
|
|
m, n = self.num_row, self.num_col
|
|
localDiff = [[] for _ in xrange(m)]
|
|
blk_sz = self.blk_sz
|
|
for r in xrange(m):
|
|
for c in xrange(n):
|
|
I_row = 0
|
|
I_col = 0
|
|
#get ssd surface
|
|
count = 0
|
|
center = self.cur_yuv[r * blk_sz:(r + 1) * blk_sz,
|
|
c * blk_sz:(c + 1) * blk_sz, 0]
|
|
ty = np.clip(r * blk_sz + int(mvs[r, c, 0]), 0, self.height - blk_sz)
|
|
tx = np.clip(c * blk_sz + int(mvs[r, c, 1]), 0, self.width - blk_sz)
|
|
target = self.ref_yuv[ty:ty + blk_sz, tx:tx + blk_sz, 0]
|
|
for y, x in {(ty - blk_sz, tx), (ty + blk_sz, tx)}:
|
|
if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
|
|
nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
|
|
I_row += np.sum(np.abs(nb - center)) - np.sum(
|
|
np.abs(target - center))
|
|
count += 1
|
|
I_row //= (count * blk_sz * blk_sz)
|
|
count = 0
|
|
for y, x in {(ty, tx - blk_sz), (ty, tx + blk_sz)}:
|
|
if 0 <= y < self.height - blk_sz and 0 <= x < self.width - blk_sz:
|
|
nb = self.ref_yuv[y:y + blk_sz, x:x + blk_sz, 0]
|
|
I_col += np.sum(np.abs(nb - center)) - np.sum(
|
|
np.abs(target - center))
|
|
count += 1
|
|
I_col //= (count * blk_sz * blk_sz)
|
|
localDiff[r].append(
|
|
np.array([[I_row * I_row, I_row * I_col],
|
|
[I_col * I_row, I_col * I_col]]))
|
|
return localDiff
|
|
|
|
"""
|
|
add smooth constraint
|
|
"""
|
|
|
|
def smooth(self, uvs, mvs):
|
|
sm_uvs = np.zeros(uvs.shape)
|
|
blk_sz = self.blk_sz
|
|
for r in xrange(self.num_row):
|
|
for c in xrange(self.num_col):
|
|
nb_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:
|
|
nb_uv += uvs[i, j] / 6.0
|
|
else:
|
|
nb_uv += uvs[r, c] / 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:
|
|
nb_uv += uvs[i, j] / 12.0
|
|
else:
|
|
nb_uv += uvs[r, c] / 12.0
|
|
mv = mvs[r, c] / blk_sz
|
|
M = self.localDiff[r][c]
|
|
P = M + self.beta * np.identity(2)
|
|
inv_P = LA.inv(P)
|
|
sm_uvs[r, c] = np.dot(inv_P, self.beta * nb_uv) + np.dot(
|
|
np.matmul(inv_P, M), mv)
|
|
return sm_uvs
|
|
|
|
def block_matching(self):
|
|
self.search.motion_field_estimation()
|
|
|
|
def motion_field_estimation(self):
|
|
#get local structure
|
|
self.localDiff = self.getRefLocalDiff(self.search.mf)
|
|
#get matching results
|
|
mvs = self.search.mf
|
|
#add smoothness constraint
|
|
uvs = mvs / self.blk_sz
|
|
for _ in xrange(self.max_iter):
|
|
uvs = self.smooth(uvs, mvs)
|
|
self.mf = uvs * self.blk_sz
|