(61d00a474) v0.9.7.1
This commit is contained in:
@@ -0,0 +1,185 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
import numpy as np
|
||||
import numpy.linalg as LA
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from scipy.sparse import csc_matrix
|
||||
from scipy.sparse.linalg import inv
|
||||
from MotionEST import MotionEST
|
||||
"""Anandan Model"""
|
||||
|
||||
|
||||
class Anandan(MotionEST):
|
||||
"""
|
||||
constructor:
|
||||
cur_f: current frame
|
||||
ref_f: reference frame
|
||||
blk_sz: block size
|
||||
beta: smooth constrain weight
|
||||
k1,k2,k3: confidence coefficients
|
||||
max_iter: maximum number of iterations
|
||||
"""
|
||||
|
||||
def __init__(self, cur_f, ref_f, blk_sz, beta, k1, k2, k3, max_iter=100):
|
||||
super(Anandan, self).__init__(cur_f, ref_f, blk_sz)
|
||||
self.levels = int(np.log2(blk_sz))
|
||||
self.intensity_hierarchy()
|
||||
self.c_maxs = []
|
||||
self.c_mins = []
|
||||
self.e_maxs = []
|
||||
self.e_mins = []
|
||||
for l in xrange(self.levels + 1):
|
||||
c_max, c_min, e_max, e_min = self.get_curvature(self.cur_Is[l])
|
||||
self.c_maxs.append(c_max)
|
||||
self.c_mins.append(c_min)
|
||||
self.e_maxs.append(e_max)
|
||||
self.e_mins.append(e_min)
|
||||
self.beta = beta
|
||||
self.k1, self.k2, self.k3 = k1, k2, k3
|
||||
self.max_iter = max_iter
|
||||
|
||||
"""
|
||||
build intensity hierarchy
|
||||
"""
|
||||
|
||||
def intensity_hierarchy(self):
|
||||
level = 0
|
||||
self.cur_Is = []
|
||||
self.ref_Is = []
|
||||
#build each level itensity by using gaussian filters
|
||||
while level <= self.levels:
|
||||
cur_I = gaussian_filter(self.cur_yuv[:, :, 0], sigma=(2**level) * 0.56)
|
||||
ref_I = gaussian_filter(self.ref_yuv[:, :, 0], sigma=(2**level) * 0.56)
|
||||
self.ref_Is.append(ref_I)
|
||||
self.cur_Is.append(cur_I)
|
||||
level += 1
|
||||
|
||||
"""
|
||||
get curvature of each block
|
||||
"""
|
||||
|
||||
def get_curvature(self, I):
|
||||
c_max = np.zeros((self.num_row, self.num_col))
|
||||
c_min = np.zeros((self.num_row, self.num_col))
|
||||
e_max = np.zeros((self.num_row, self.num_col, 2))
|
||||
e_min = np.zeros((self.num_row, self.num_col, 2))
|
||||
for r in xrange(self.num_row):
|
||||
for c in xrange(self.num_col):
|
||||
h11, h12, h21, h22 = 0, 0, 0, 0
|
||||
for i in xrange(r * self.blk_sz, r * self.blk_sz + self.blk_sz):
|
||||
for j in xrange(c * self.blk_sz, c * self.blk_sz + self.blk_sz):
|
||||
if 0 <= i < self.height - 1 and 0 <= j < self.width - 1:
|
||||
Ix = I[i][j + 1] - I[i][j]
|
||||
Iy = I[i + 1][j] - I[i][j]
|
||||
h11 += Iy * Iy
|
||||
h12 += Ix * Iy
|
||||
h21 += Ix * Iy
|
||||
h22 += Ix * Ix
|
||||
U, S, _ = LA.svd(np.array([[h11, h12], [h21, h22]]))
|
||||
c_max[r, c], c_min[r, c] = S[0], S[1]
|
||||
e_max[r, c] = U[:, 0]
|
||||
e_min[r, c] = U[:, 1]
|
||||
return c_max, c_min, e_max, e_min
|
||||
|
||||
"""
|
||||
get ssd of motion vector:
|
||||
cur_I: current intensity
|
||||
ref_I: reference intensity
|
||||
center: current position
|
||||
mv: motion vector
|
||||
"""
|
||||
|
||||
def get_ssd(self, cur_I, ref_I, center, mv):
|
||||
ssd = 0
|
||||
for r in xrange(int(center[0]), int(center[0]) + self.blk_sz):
|
||||
for c in xrange(int(center[1]), int(center[1]) + self.blk_sz):
|
||||
if 0 <= r < self.height and 0 <= c < self.width:
|
||||
tr, tc = r + int(mv[0]), c + int(mv[1])
|
||||
if 0 <= tr < self.height and 0 <= tc < self.width:
|
||||
ssd += (ref_I[tr, tc] - cur_I[r, c])**2
|
||||
else:
|
||||
ssd += cur_I[r, c]**2
|
||||
return ssd
|
||||
|
||||
"""
|
||||
get region match of level l
|
||||
l: current level
|
||||
last_mvs: matchine results of last level
|
||||
radius: movenment radius
|
||||
"""
|
||||
|
||||
def region_match(self, l, last_mvs, radius):
|
||||
mvs = np.zeros((self.num_row, self.num_col, 2))
|
||||
min_ssds = np.zeros((self.num_row, self.num_col))
|
||||
for r in xrange(self.num_row):
|
||||
for c in xrange(self.num_col):
|
||||
center = np.array([r * self.blk_sz, c * self.blk_sz])
|
||||
#use overlap hierarchy policy
|
||||
init_mvs = []
|
||||
if last_mvs is None:
|
||||
init_mvs = [np.array([0, 0])]
|
||||
else:
|
||||
for i, j in {(r, c), (r, c + 1), (r + 1, c), (r + 1, c + 1)}:
|
||||
if 0 <= i < last_mvs.shape[0] and 0 <= j < last_mvs.shape[1]:
|
||||
init_mvs.append(last_mvs[i, j])
|
||||
#use last matching results as the start postion as current level
|
||||
min_ssd = None
|
||||
min_mv = None
|
||||
for init_mv in init_mvs:
|
||||
for i in xrange(-2, 3):
|
||||
for j in xrange(-2, 3):
|
||||
mv = init_mv + np.array([i, j]) * radius
|
||||
ssd = self.get_ssd(self.cur_Is[l], self.ref_Is[l], center, mv)
|
||||
if min_ssd is None or ssd < min_ssd:
|
||||
min_ssd = ssd
|
||||
min_mv = mv
|
||||
min_ssds[r, c] = min_ssd
|
||||
mvs[r, c] = min_mv
|
||||
return mvs, min_ssds
|
||||
|
||||
"""
|
||||
smooth motion field based on neighbor constraint
|
||||
uvs: current estimation
|
||||
mvs: matching results
|
||||
min_ssds: minimum ssd of matching results
|
||||
l: current level
|
||||
"""
|
||||
|
||||
def smooth(self, uvs, mvs, min_ssds, l):
|
||||
sm_uvs = np.zeros((self.num_row, self.num_col, 2))
|
||||
c_max = self.c_maxs[l]
|
||||
c_min = self.c_mins[l]
|
||||
e_max = self.e_maxs[l]
|
||||
e_min = self.e_mins[l]
|
||||
for r in xrange(self.num_row):
|
||||
for c in xrange(self.num_col):
|
||||
w_max = c_max[r, c] / (
|
||||
self.k1 + self.k2 * min_ssds[r, c] + self.k3 * c_max[r, c])
|
||||
w_min = c_min[r, c] / (
|
||||
self.k1 + self.k2 * min_ssds[r, c] + self.k3 * c_min[r, c])
|
||||
w = w_max * w_min / (w_max + w_min + 1e-6)
|
||||
if w < 0:
|
||||
w = 0
|
||||
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 += 0.25 * uvs[i, j]
|
||||
sm_uvs[r, c] = (w * w * mvs[r, c] + self.beta * avg_uv) / (
|
||||
self.beta + w * w)
|
||||
return sm_uvs
|
||||
|
||||
"""
|
||||
motion field estimation
|
||||
"""
|
||||
|
||||
def motion_field_estimation(self):
|
||||
last_mvs = None
|
||||
for l in xrange(self.levels, -1, -1):
|
||||
mvs, min_ssds = self.region_match(l, last_mvs, 2**l)
|
||||
uvs = np.zeros(mvs.shape)
|
||||
for _ in xrange(self.max_iter):
|
||||
uvs = self.smooth(uvs, mvs, min_ssds, l)
|
||||
last_mvs = uvs
|
||||
for r in xrange(self.num_row):
|
||||
for c in xrange(self.num_col):
|
||||
self.mf[r, c] = uvs[r, c]
|
||||
@@ -0,0 +1,251 @@
|
||||
#!/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
|
||||
@@ -0,0 +1,40 @@
|
||||
#!/ 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
|
||||
"""
|
||||
|
||||
|
||||
class GroundTruth(MotionEST):
|
||||
"""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, 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)
|
||||
if gt_path:
|
||||
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
|
||||
@@ -0,0 +1,204 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
import numpy as np
|
||||
import numpy.linalg as LA
|
||||
from scipy.ndimage.filters import gaussian_filter
|
||||
from scipy.sparse import csc_matrix
|
||||
from scipy.sparse.linalg import inv
|
||||
from MotionEST import MotionEST
|
||||
"""Horn & Schunck Model"""
|
||||
|
||||
|
||||
class HornSchunck(MotionEST):
|
||||
"""
|
||||
constructor:
|
||||
cur_f: current frame
|
||||
ref_f: reference frame
|
||||
blk_sz: block size
|
||||
alpha: smooth constrain weight
|
||||
sigma: gaussian blur parameter
|
||||
"""
|
||||
|
||||
def __init__(self, cur_f, ref_f, blk_sz, alpha, sigma, max_iter=100):
|
||||
super(HornSchunck, self).__init__(cur_f, ref_f, blk_sz)
|
||||
self.cur_I, self.ref_I = self.getIntensity()
|
||||
#perform gaussian blur to smooth the intensity
|
||||
self.cur_I = gaussian_filter(self.cur_I, sigma=sigma)
|
||||
self.ref_I = gaussian_filter(self.ref_I, sigma=sigma)
|
||||
self.alpha = alpha
|
||||
self.max_iter = max_iter
|
||||
self.Ix, self.Iy, self.It = self.intensityDiff()
|
||||
|
||||
"""
|
||||
Build Frame Intensity
|
||||
"""
|
||||
|
||||
def getIntensity(self):
|
||||
cur_I = np.zeros((self.num_row, self.num_col))
|
||||
ref_I = np.zeros((self.num_row, self.num_col))
|
||||
#use average intensity as block's intensity
|
||||
for i in xrange(self.num_row):
|
||||
for j in xrange(self.num_col):
|
||||
r = i * self.blk_sz
|
||||
c = j * self.blk_sz
|
||||
cur_I[i, j] = np.mean(self.cur_yuv[r:r + self.blk_sz, c:c + self.blk_sz,
|
||||
0])
|
||||
ref_I[i, j] = np.mean(self.ref_yuv[r:r + self.blk_sz, c:c + self.blk_sz,
|
||||
0])
|
||||
return cur_I, ref_I
|
||||
|
||||
"""
|
||||
Get First Order Derivative
|
||||
"""
|
||||
|
||||
def intensityDiff(self):
|
||||
Ix = np.zeros((self.num_row, self.num_col))
|
||||
Iy = np.zeros((self.num_row, self.num_col))
|
||||
It = np.zeros((self.num_row, self.num_col))
|
||||
sz = self.blk_sz
|
||||
for i in xrange(self.num_row - 1):
|
||||
for j in xrange(self.num_col - 1):
|
||||
"""
|
||||
Ix:
|
||||
(i ,j) <--- (i ,j+1)
|
||||
(i+1,j) <--- (i+1,j+1)
|
||||
"""
|
||||
count = 0
|
||||
for r, c in {(i, j + 1), (i + 1, j + 1)}:
|
||||
if 0 <= r < self.num_row and 0 < c < self.num_col:
|
||||
Ix[i, j] += (
|
||||
self.cur_I[r, c] - self.cur_I[r, c - 1] + self.ref_I[r, c] -
|
||||
self.ref_I[r, c - 1])
|
||||
count += 2
|
||||
Ix[i, j] /= count
|
||||
"""
|
||||
Iy:
|
||||
(i ,j) (i ,j+1)
|
||||
^ ^
|
||||
| |
|
||||
(i+1,j) (i+1,j+1)
|
||||
"""
|
||||
count = 0
|
||||
for r, c in {(i + 1, j), (i + 1, j + 1)}:
|
||||
if 0 < r < self.num_row and 0 <= c < self.num_col:
|
||||
Iy[i, j] += (
|
||||
self.cur_I[r, c] - self.cur_I[r - 1, c] + self.ref_I[r, c] -
|
||||
self.ref_I[r - 1, c])
|
||||
count += 2
|
||||
Iy[i, j] /= count
|
||||
count = 0
|
||||
#It:
|
||||
for r in xrange(i, i + 2):
|
||||
for c in xrange(j, j + 2):
|
||||
if 0 <= r < self.num_row and 0 <= c < self.num_col:
|
||||
It[i, j] += (self.ref_I[r, c] - self.cur_I[r, c])
|
||||
count += 1
|
||||
It[i, j] /= count
|
||||
return Ix, Iy, It
|
||||
|
||||
"""
|
||||
Get weighted average of neighbor motion vectors
|
||||
for evaluation of laplacian
|
||||
"""
|
||||
|
||||
def averageMV(self):
|
||||
avg = np.zeros((self.num_row, self.num_col, 2))
|
||||
"""
|
||||
1/12 --- 1/6 --- 1/12
|
||||
| | |
|
||||
1/6 --- -1/8 --- 1/6
|
||||
| | |
|
||||
1/12 --- 1/6 --- 1/12
|
||||
"""
|
||||
for i in xrange(self.num_row):
|
||||
for j in xrange(self.num_col):
|
||||
for r, c in {(-1, 0), (1, 0), (0, -1), (0, 1)}:
|
||||
if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
|
||||
avg[i, j] += self.mf[i + r, j + c] / 6.0
|
||||
for r, c in {(-1, -1), (-1, 1), (1, -1), (1, 1)}:
|
||||
if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
|
||||
avg[i, j] += self.mf[i + r, j + c] / 12.0
|
||||
return avg
|
||||
|
||||
def motion_field_estimation(self):
|
||||
count = 0
|
||||
"""
|
||||
u_{n+1} = ~u_n - Ix(Ix.~u_n+Iy.~v+It)/(IxIx+IyIy+alpha^2)
|
||||
v_{n+1} = ~v_n - Iy(Ix.~u_n+Iy.~v+It)/(IxIx+IyIy+alpha^2)
|
||||
"""
|
||||
denom = self.alpha**2 + np.power(self.Ix, 2) + np.power(self.Iy, 2)
|
||||
while count < self.max_iter:
|
||||
avg = self.averageMV()
|
||||
self.mf[:, :, 1] = avg[:, :, 1] - self.Ix * (
|
||||
self.Ix * avg[:, :, 1] + self.Iy * avg[:, :, 0] + self.It) / denom
|
||||
self.mf[:, :, 0] = avg[:, :, 0] - self.Iy * (
|
||||
self.Ix * avg[:, :, 1] + self.Iy * avg[:, :, 0] + self.It) / denom
|
||||
count += 1
|
||||
self.mf *= self.blk_sz
|
||||
|
||||
def motion_field_estimation_mat(self):
|
||||
row_idx = []
|
||||
col_idx = []
|
||||
data = []
|
||||
|
||||
N = 2 * self.num_row * self.num_col
|
||||
b = np.zeros((N, 1))
|
||||
for i in xrange(self.num_row):
|
||||
for j in xrange(self.num_col):
|
||||
"""(IxIx+alpha^2)u+IxIy.v-alpha^2~u IxIy.u+(IyIy+alpha^2)v-alpha^2~v"""
|
||||
u_idx = i * 2 * self.num_col + 2 * j
|
||||
v_idx = u_idx + 1
|
||||
b[u_idx, 0] = -self.Ix[i, j] * self.It[i, j]
|
||||
b[v_idx, 0] = -self.Iy[i, j] * self.It[i, j]
|
||||
#u: (IxIx+alpha^2)u
|
||||
row_idx.append(u_idx)
|
||||
col_idx.append(u_idx)
|
||||
data.append(self.Ix[i, j] * self.Ix[i, j] + self.alpha**2)
|
||||
#IxIy.v
|
||||
row_idx.append(u_idx)
|
||||
col_idx.append(v_idx)
|
||||
data.append(self.Ix[i, j] * self.Iy[i, j])
|
||||
|
||||
#v: IxIy.u
|
||||
row_idx.append(v_idx)
|
||||
col_idx.append(u_idx)
|
||||
data.append(self.Ix[i, j] * self.Iy[i, j])
|
||||
#(IyIy+alpha^2)v
|
||||
row_idx.append(v_idx)
|
||||
col_idx.append(v_idx)
|
||||
data.append(self.Iy[i, j] * self.Iy[i, j] + self.alpha**2)
|
||||
|
||||
#-alpha^2~u
|
||||
#-alpha^2~v
|
||||
for r, c in {(-1, 0), (1, 0), (0, -1), (0, 1)}:
|
||||
if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
|
||||
u_nb = (i + r) * 2 * self.num_col + 2 * (j + c)
|
||||
v_nb = u_nb + 1
|
||||
|
||||
row_idx.append(u_idx)
|
||||
col_idx.append(u_nb)
|
||||
data.append(-1 * self.alpha**2 / 6.0)
|
||||
|
||||
row_idx.append(v_idx)
|
||||
col_idx.append(v_nb)
|
||||
data.append(-1 * self.alpha**2 / 6.0)
|
||||
for r, c in {(-1, -1), (-1, 1), (1, -1), (1, 1)}:
|
||||
if 0 <= i + r < self.num_row and 0 <= j + c < self.num_col:
|
||||
u_nb = (i + r) * 2 * self.num_col + 2 * (j + c)
|
||||
v_nb = u_nb + 1
|
||||
|
||||
row_idx.append(u_idx)
|
||||
col_idx.append(u_nb)
|
||||
data.append(-1 * self.alpha**2 / 12.0)
|
||||
|
||||
row_idx.append(v_idx)
|
||||
col_idx.append(v_nb)
|
||||
data.append(-1 * self.alpha**2 / 12.0)
|
||||
M = csc_matrix((data, (row_idx, col_idx)), shape=(N, N))
|
||||
M_inv = inv(M)
|
||||
uv = M_inv.dot(b)
|
||||
|
||||
for i in xrange(self.num_row):
|
||||
for j in xrange(self.num_col):
|
||||
self.mf[i, j, 0] = uv[i * 2 * self.num_col + 2 * j + 1, 0] * self.blk_sz
|
||||
self.mf[i, j, 1] = uv[i * 2 * self.num_col + 2 * j, 0] * self.blk_sz
|
||||
@@ -0,0 +1,109 @@
|
||||
#!/ usr / bin / env python
|
||||
#coding : utf - 8
|
||||
import numpy as np
|
||||
import numpy.linalg as LA
|
||||
import matplotlib.pyplot as plt
|
||||
from Util import drawMF, MSE
|
||||
"""The Base Class of Estimators"""
|
||||
|
||||
|
||||
class MotionEST(object):
|
||||
"""
|
||||
constructor:
|
||||
cur_f: current frame
|
||||
ref_f: reference frame
|
||||
blk_sz: block size
|
||||
"""
|
||||
|
||||
def __init__(self, cur_f, ref_f, blk_sz):
|
||||
self.cur_f = cur_f
|
||||
self.ref_f = ref_f
|
||||
self.blk_sz = blk_sz
|
||||
#convert RGB to YUV
|
||||
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]
|
||||
#motion field size
|
||||
self.num_row = self.height // self.blk_sz
|
||||
self.num_col = self.width // self.blk_sz
|
||||
#initialize motion field
|
||||
self.mf = np.zeros((self.num_row, self.num_col, 2))
|
||||
|
||||
"""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
|
||||
"""
|
||||
|
||||
def block_dist(self, cur_r, cur_c, mv, metric=MSE):
|
||||
cur_x = cur_c * self.blk_sz
|
||||
cur_y = cur_r * self.blk_sz
|
||||
h = min(self.blk_sz, self.height - cur_y)
|
||||
w = min(self.blk_sz, self.width - cur_x)
|
||||
cur_blk = self.cur_yuv[cur_y:cur_y + h, cur_x:cur_x + w, :]
|
||||
ref_x = int(cur_x + mv[1])
|
||||
ref_y = int(cur_y + mv[0])
|
||||
if 0 <= ref_x < self.width - w and 0 <= ref_y < self.height - h:
|
||||
ref_blk = self.ref_yuv[ref_y:ref_y + h, ref_x:ref_x + w, :]
|
||||
else:
|
||||
ref_blk = np.zeros((h, w, 3))
|
||||
return metric(cur_blk, ref_blk)
|
||||
|
||||
"""
|
||||
distortion of motion field
|
||||
"""
|
||||
|
||||
def distortion(self, mask=None, metric=MSE):
|
||||
loss = 0
|
||||
count = 0
|
||||
for i in xrange(self.num_row):
|
||||
for j in xrange(self.num_col):
|
||||
if mask is not None and mask[i, j]:
|
||||
continue
|
||||
loss += self.block_dist(i, j, self.mf[i, j], metric)
|
||||
count += 1
|
||||
return loss / count
|
||||
|
||||
"""evaluation compare the difference with ground truth"""
|
||||
|
||||
def motion_field_evaluation(self, ground_truth):
|
||||
loss = 0
|
||||
count = 0
|
||||
gt = ground_truth.mf
|
||||
mask = ground_truth.mask
|
||||
for i in xrange(self.num_row):
|
||||
for j in xrange(self.num_col):
|
||||
if mask is not None and mask[i][j]:
|
||||
continue
|
||||
loss += LA.norm(gt[i, j] - self.mf[i, j])
|
||||
count += 1
|
||||
return loss / count
|
||||
|
||||
"""render the motion field"""
|
||||
|
||||
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 * size, size * self.height / self.width))
|
||||
plt.subplot(1, n_row, 1)
|
||||
plt.imshow(cur_mf)
|
||||
plt.title('Estimated Motion Field')
|
||||
if ground_truth is not None:
|
||||
plt.subplot(1, n_row, 2)
|
||||
plt.imshow(gt_mf)
|
||||
plt.title('Ground Truth')
|
||||
plt.tight_layout()
|
||||
plt.show()
|
||||
@@ -0,0 +1,213 @@
|
||||
#!/usr/bin/env python
|
||||
# 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
|
||||
@@ -0,0 +1,38 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
import numpy as np
|
||||
import numpy.linalg as LA
|
||||
import matplotlib.pyplot as plt
|
||||
from scipy.ndimage import filters
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
def MSE(blk1, blk2):
|
||||
return np.mean(
|
||||
LA.norm(
|
||||
np.array(blk1, dtype=np.int) - np.array(blk2, dtype=np.int), axis=2))
|
||||
|
||||
|
||||
def drawMF(img, blk_sz, mf):
|
||||
img_rgba = img.convert('RGBA')
|
||||
mf_layer = Image.new(mode='RGBA', size=img_rgba.size, color=(0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(mf_layer)
|
||||
width = img_rgba.size[0]
|
||||
height = img_rgba.size[1]
|
||||
num_row = height // blk_sz
|
||||
num_col = width // blk_sz
|
||||
for i in xrange(num_row):
|
||||
left = (0, i * blk_sz)
|
||||
right = (width, i * blk_sz)
|
||||
draw.line([left, right], fill=(0, 0, 255, 255))
|
||||
for j in xrange(num_col):
|
||||
up = (j * blk_sz, 0)
|
||||
down = (j * blk_sz, height)
|
||||
draw.line([up, down], fill=(0, 0, 255, 255))
|
||||
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)."""
|
||||
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)
|
||||
@@ -0,0 +1,76 @@
|
||||
import argparse
|
||||
from os import listdir, path
|
||||
from PIL import Image
|
||||
import sys
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--frame_path", default="../data/frame/", type=str)
|
||||
parser.add_argument("--frame_rate", default="25:1", type=str)
|
||||
parser.add_argument("--interlacing", default="Ip", type=str)
|
||||
parser.add_argument("--pix_ratio", default="0:0", type=str)
|
||||
parser.add_argument("--color_space", default="4:2:0", type=str)
|
||||
parser.add_argument("--output", default="output.y4m", type=str)
|
||||
|
||||
|
||||
def generate(args, frames):
|
||||
if len(frames) == 0:
|
||||
return
|
||||
#sort the frames based on the frame index
|
||||
frames = sorted(frames, key=lambda x: x[0])
|
||||
#convert the frames to YUV form
|
||||
frames = [f.convert("YCbCr") for _, f in frames]
|
||||
#write the header
|
||||
header = "YUV4MPEG2 W%d H%d F%s %s A%s" % (frames[0].width, frames[0].height,
|
||||
args.frame_rate, args.interlacing,
|
||||
args.pix_ratio)
|
||||
cs = args.color_space.split(":")
|
||||
header += " C%s%s%s\n" % (cs[0], cs[1], cs[2])
|
||||
#estimate the sample step based on subsample value
|
||||
subsamples = [int(c) for c in cs]
|
||||
r_step = [1, int(subsamples[2] == 0) + 1, int(subsamples[2] == 0) + 1]
|
||||
c_step = [1, 4 // subsamples[1], 4 // subsamples[1]]
|
||||
#write in frames
|
||||
with open(args.output, "wb") as y4m:
|
||||
y4m.write(header)
|
||||
for f in frames:
|
||||
y4m.write("FRAME\n")
|
||||
px = f.load()
|
||||
for k in xrange(3):
|
||||
for i in xrange(0, f.height, r_step[k]):
|
||||
for j in xrange(0, f.width, c_step[k]):
|
||||
yuv = px[j, i]
|
||||
y4m.write(chr(yuv[k]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parser.parse_args()
|
||||
frames = []
|
||||
frames_mv = []
|
||||
for filename in listdir(args.frame_path):
|
||||
name, ext = filename.split(".")
|
||||
if ext == "png":
|
||||
name_parse = name.split("_")
|
||||
idx = int(name_parse[-1])
|
||||
img = Image.open(path.join(args.frame_path, filename))
|
||||
if name_parse[-2] == "mv":
|
||||
frames_mv.append((idx, img))
|
||||
else:
|
||||
frames.append((idx, img))
|
||||
if len(frames) == 0:
|
||||
print("No frames in directory: " + args.frame_path)
|
||||
sys.exit()
|
||||
print("----------------------Y4M Info----------------------")
|
||||
print("width: %d" % frames[0][1].width)
|
||||
print("height: %d" % frames[0][1].height)
|
||||
print("#frame: %d" % len(frames))
|
||||
print("frame rate: %s" % args.frame_rate)
|
||||
print("interlacing: %s" % args.interlacing)
|
||||
print("pixel ratio: %s" % args.pix_ratio)
|
||||
print("color space: %s" % args.color_space)
|
||||
print("----------------------------------------------------")
|
||||
|
||||
print("Generating ...")
|
||||
generate(args, frames)
|
||||
if len(frames_mv) != 0:
|
||||
args.output = args.output.replace(".y4m", "_mv.y4m")
|
||||
generate(args, frames_mv)
|
||||
+163
@@ -0,0 +1,163 @@
|
||||
/*
|
||||
*AABB bounding box
|
||||
*Bouding Volume Hierarchy
|
||||
*/
|
||||
class BoundingBox {
|
||||
float min_x, min_y, min_z, max_x, max_y, max_z;
|
||||
PVector center;
|
||||
BoundingBox() {
|
||||
min_x = Float.POSITIVE_INFINITY;
|
||||
min_y = Float.POSITIVE_INFINITY;
|
||||
min_z = Float.POSITIVE_INFINITY;
|
||||
max_x = Float.NEGATIVE_INFINITY;
|
||||
max_y = Float.NEGATIVE_INFINITY;
|
||||
max_z = Float.NEGATIVE_INFINITY;
|
||||
center = new PVector();
|
||||
}
|
||||
// build a bounding box for a triangle
|
||||
void create(Triangle t) {
|
||||
min_x = min(t.p1.x, min(t.p2.x, t.p3.x));
|
||||
max_x = max(t.p1.x, max(t.p2.x, t.p3.x));
|
||||
|
||||
min_y = min(t.p1.y, min(t.p2.y, t.p3.y));
|
||||
max_y = max(t.p1.y, max(t.p2.y, t.p3.y));
|
||||
|
||||
min_z = min(t.p1.z, min(t.p2.z, t.p3.z));
|
||||
max_z = max(t.p1.z, max(t.p2.z, t.p3.z));
|
||||
center.x = (max_x + min_x) / 2;
|
||||
center.y = (max_y + min_y) / 2;
|
||||
center.z = (max_z + min_z) / 2;
|
||||
}
|
||||
// merge two bounding boxes
|
||||
void add(BoundingBox bbx) {
|
||||
min_x = min(min_x, bbx.min_x);
|
||||
min_y = min(min_y, bbx.min_y);
|
||||
min_z = min(min_z, bbx.min_z);
|
||||
|
||||
max_x = max(max_x, bbx.max_x);
|
||||
max_y = max(max_y, bbx.max_y);
|
||||
max_z = max(max_z, bbx.max_z);
|
||||
center.x = (max_x + min_x) / 2;
|
||||
center.y = (max_y + min_y) / 2;
|
||||
center.z = (max_z + min_z) / 2;
|
||||
}
|
||||
// get bounding box center axis value
|
||||
float getCenterAxisValue(int axis) {
|
||||
if (axis == 1) {
|
||||
return center.x;
|
||||
} else if (axis == 2) {
|
||||
return center.y;
|
||||
}
|
||||
// when axis == 3
|
||||
return center.z;
|
||||
}
|
||||
// check if a ray is intersected with the bounding box
|
||||
boolean intersect(Ray r) {
|
||||
float tmin, tmax;
|
||||
if (r.dir.x >= 0) {
|
||||
tmin = (min_x - r.ori.x) * (1.0f / r.dir.x);
|
||||
tmax = (max_x - r.ori.x) * (1.0f / r.dir.x);
|
||||
} else {
|
||||
tmin = (max_x - r.ori.x) * (1.0f / r.dir.x);
|
||||
tmax = (min_x - r.ori.x) * (1.0f / r.dir.x);
|
||||
}
|
||||
|
||||
float tymin, tymax;
|
||||
if (r.dir.y >= 0) {
|
||||
tymin = (min_y - r.ori.y) * (1.0f / r.dir.y);
|
||||
tymax = (max_y - r.ori.y) * (1.0f / r.dir.y);
|
||||
} else {
|
||||
tymin = (max_y - r.ori.y) * (1.0f / r.dir.y);
|
||||
tymax = (min_y - r.ori.y) * (1.0f / r.dir.y);
|
||||
}
|
||||
|
||||
if (tmax < tymin || tymax < tmin) {
|
||||
return false;
|
||||
}
|
||||
|
||||
tmin = tmin < tymin ? tymin : tmin;
|
||||
tmax = tmax > tymax ? tymax : tmax;
|
||||
|
||||
float tzmin, tzmax;
|
||||
if (r.dir.z >= 0) {
|
||||
tzmin = (min_z - r.ori.z) * (1.0f / r.dir.z);
|
||||
tzmax = (max_z - r.ori.z) * (1.0f / r.dir.z);
|
||||
} else {
|
||||
tzmin = (max_z - r.ori.z) * (1.0f / r.dir.z);
|
||||
tzmax = (min_z - r.ori.z) * (1.0f / r.dir.z);
|
||||
}
|
||||
if (tmax < tzmin || tmin > tzmax) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
}
|
||||
// Bounding Volume Hierarchy
|
||||
class BVH {
|
||||
// Binary Tree
|
||||
BVH left, right;
|
||||
BoundingBox overall_bbx;
|
||||
ArrayList<Triangle> mesh;
|
||||
BVH(ArrayList<Triangle> mesh) {
|
||||
this.mesh = mesh;
|
||||
overall_bbx = new BoundingBox();
|
||||
left = null;
|
||||
right = null;
|
||||
int mesh_size = this.mesh.size();
|
||||
if (mesh_size <= 1) {
|
||||
return;
|
||||
}
|
||||
// random select an axis
|
||||
int axis = int(random(100)) % 3 + 1;
|
||||
// build bounding box and save the selected center component
|
||||
float[] axis_values = new float[mesh_size];
|
||||
for (int i = 0; i < mesh_size; i++) {
|
||||
Triangle t = this.mesh.get(i);
|
||||
overall_bbx.add(t.bbx);
|
||||
axis_values[i] = t.bbx.getCenterAxisValue(axis);
|
||||
}
|
||||
// find the median value of selected center component as pivot
|
||||
axis_values = sort(axis_values);
|
||||
float pivot;
|
||||
if (mesh_size % 2 == 1) {
|
||||
pivot = axis_values[mesh_size / 2];
|
||||
} else {
|
||||
pivot =
|
||||
0.5f * (axis_values[mesh_size / 2 - 1] + axis_values[mesh_size / 2]);
|
||||
}
|
||||
// Build left node and right node by partitioning the mesh based on triangle
|
||||
// bounding box center component value
|
||||
ArrayList<Triangle> left_mesh = new ArrayList<Triangle>();
|
||||
ArrayList<Triangle> right_mesh = new ArrayList<Triangle>();
|
||||
for (int i = 0; i < mesh_size; i++) {
|
||||
Triangle t = this.mesh.get(i);
|
||||
if (t.bbx.getCenterAxisValue(axis) < pivot) {
|
||||
left_mesh.add(t);
|
||||
} else if (t.bbx.getCenterAxisValue(axis) > pivot) {
|
||||
right_mesh.add(t);
|
||||
} else if (left_mesh.size() < right_mesh.size()) {
|
||||
left_mesh.add(t);
|
||||
} else {
|
||||
right_mesh.add(t);
|
||||
}
|
||||
}
|
||||
left = new BVH(left_mesh);
|
||||
right = new BVH(right_mesh);
|
||||
}
|
||||
// check if a ray intersect with current volume
|
||||
boolean intersect(Ray r, float[] param) {
|
||||
if (mesh.size() == 0) {
|
||||
return false;
|
||||
}
|
||||
if (mesh.size() == 1) {
|
||||
Triangle t = mesh.get(0);
|
||||
return t.intersect(r, param);
|
||||
}
|
||||
if (!overall_bbx.intersect(r)) {
|
||||
return false;
|
||||
}
|
||||
boolean left_res = left.intersect(r, param);
|
||||
boolean right_res = right.intersect(r, param);
|
||||
return left_res || right_res;
|
||||
}
|
||||
}
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
class Camera {
|
||||
// camera's field of view
|
||||
float fov;
|
||||
// camera's position, look at point and axis
|
||||
PVector pos, center, axis;
|
||||
PVector init_pos, init_center, init_axis;
|
||||
float move_speed;
|
||||
float rot_speed;
|
||||
Camera(float fov, PVector pos, PVector center, PVector axis) {
|
||||
this.fov = fov;
|
||||
this.pos = pos;
|
||||
this.center = center;
|
||||
this.axis = axis;
|
||||
this.axis.normalize();
|
||||
move_speed = 0.001;
|
||||
rot_speed = 0.01 * PI;
|
||||
init_pos = pos.copy();
|
||||
init_center = center.copy();
|
||||
init_axis = axis.copy();
|
||||
}
|
||||
|
||||
Camera copy() {
|
||||
Camera cam = new Camera(fov, pos.copy(), center.copy(), axis.copy());
|
||||
return cam;
|
||||
}
|
||||
|
||||
PVector project(PVector pos) {
|
||||
PVector proj = MatxVec3(getCameraMat(), PVector.sub(pos, this.pos));
|
||||
proj.x = (float)height / 2.0 * proj.x / proj.z / tan(fov / 2.0f);
|
||||
proj.y = (float)height / 2.0 * proj.y / proj.z / tan(fov / 2.0f);
|
||||
proj.z = proj.z;
|
||||
return proj;
|
||||
}
|
||||
|
||||
float[] getCameraMat() {
|
||||
float[] mat = new float[9];
|
||||
PVector dir = PVector.sub(center, pos);
|
||||
dir.normalize();
|
||||
PVector left = dir.cross(axis);
|
||||
left.normalize();
|
||||
// processing camera system does not follow right hand rule
|
||||
mat[0] = -left.x;
|
||||
mat[1] = -left.y;
|
||||
mat[2] = -left.z;
|
||||
mat[3] = axis.x;
|
||||
mat[4] = axis.y;
|
||||
mat[5] = axis.z;
|
||||
mat[6] = dir.x;
|
||||
mat[7] = dir.y;
|
||||
mat[8] = dir.z;
|
||||
|
||||
return mat;
|
||||
}
|
||||
|
||||
void run() {
|
||||
PVector dir, left;
|
||||
if (mousePressed) {
|
||||
float angleX = (float)mouseX / width * PI - PI / 2;
|
||||
float angleY = (float)mouseY / height * PI - PI;
|
||||
PVector diff = PVector.sub(center, pos);
|
||||
float radius = diff.mag();
|
||||
pos.x = radius * sin(angleY) * sin(angleX) + center.x;
|
||||
pos.y = radius * cos(angleY) + center.y;
|
||||
pos.z = radius * sin(angleY) * cos(angleX) + center.z;
|
||||
dir = PVector.sub(center, pos);
|
||||
dir.normalize();
|
||||
PVector up = new PVector(0, 1, 0);
|
||||
left = up.cross(dir);
|
||||
left.normalize();
|
||||
axis = dir.cross(left);
|
||||
axis.normalize();
|
||||
}
|
||||
|
||||
if (keyPressed) {
|
||||
switch (key) {
|
||||
case 'w':
|
||||
dir = PVector.sub(center, pos);
|
||||
dir.normalize();
|
||||
pos = PVector.add(pos, PVector.mult(dir, move_speed));
|
||||
center = PVector.add(center, PVector.mult(dir, move_speed));
|
||||
break;
|
||||
case 's':
|
||||
dir = PVector.sub(center, pos);
|
||||
dir.normalize();
|
||||
pos = PVector.sub(pos, PVector.mult(dir, move_speed));
|
||||
center = PVector.sub(center, PVector.mult(dir, move_speed));
|
||||
break;
|
||||
case 'a':
|
||||
dir = PVector.sub(center, pos);
|
||||
dir.normalize();
|
||||
left = axis.cross(dir);
|
||||
left.normalize();
|
||||
pos = PVector.add(pos, PVector.mult(left, move_speed));
|
||||
center = PVector.add(center, PVector.mult(left, move_speed));
|
||||
break;
|
||||
case 'd':
|
||||
dir = PVector.sub(center, pos);
|
||||
dir.normalize();
|
||||
left = axis.cross(dir);
|
||||
left.normalize();
|
||||
pos = PVector.sub(pos, PVector.mult(left, move_speed));
|
||||
center = PVector.sub(center, PVector.mult(left, move_speed));
|
||||
break;
|
||||
case 'r':
|
||||
dir = PVector.sub(center, pos);
|
||||
dir.normalize();
|
||||
float[] mat = getRotationMat3x3(rot_speed, dir.x, dir.y, dir.z);
|
||||
axis = MatxVec3(mat, axis);
|
||||
axis.normalize();
|
||||
break;
|
||||
case 'b':
|
||||
pos = init_pos.copy();
|
||||
center = init_center.copy();
|
||||
axis = init_axis.copy();
|
||||
break;
|
||||
case '+': move_speed *= 2.0f; break;
|
||||
case '-': move_speed /= 2.0; break;
|
||||
case CODED:
|
||||
if (keyCode == UP) {
|
||||
pos = PVector.add(pos, PVector.mult(axis, move_speed));
|
||||
center = PVector.add(center, PVector.mult(axis, move_speed));
|
||||
} else if (keyCode == DOWN) {
|
||||
pos = PVector.sub(pos, PVector.mult(axis, move_speed));
|
||||
center = PVector.sub(center, PVector.mult(axis, move_speed));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
void open() {
|
||||
perspective(fov, float(width) / height, 1e-6, 1e5);
|
||||
camera(pos.x, pos.y, pos.z, center.x, center.y, center.z, axis.x, axis.y,
|
||||
axis.z);
|
||||
}
|
||||
void close() {
|
||||
ortho(-width, 0, -height, 0);
|
||||
camera(0, 0, 0, 0, 0, 1, 0, 1, 0);
|
||||
}
|
||||
}
|
||||
+102
@@ -0,0 +1,102 @@
|
||||
class MotionField {
|
||||
int block_size;
|
||||
ArrayList<PVector> motion_field;
|
||||
MotionField(int block_size) {
|
||||
this.block_size = block_size;
|
||||
motion_field = new ArrayList<PVector>();
|
||||
}
|
||||
|
||||
void update(Camera last_cam, Camera current_cam, PointCloud point_cloud,
|
||||
BVH bvh) {
|
||||
// clear motion field
|
||||
motion_field = new ArrayList<PVector>();
|
||||
int r_num = height / block_size, c_num = width / block_size;
|
||||
for (int i = 0; i < r_num * c_num; i++)
|
||||
motion_field.add(new PVector(0, 0, 0));
|
||||
// estimate motion vector of each point in point cloud
|
||||
for (int i = 0; i < point_cloud.size(); i++) {
|
||||
PVector p = point_cloud.getPosition(i);
|
||||
PVector p0 = current_cam.project(p);
|
||||
PVector p1 = last_cam.project(p);
|
||||
int row = int((p0.y + height / 2.0f) / block_size);
|
||||
int col = int((p0.x + width / 2.0f) / block_size);
|
||||
if (row >= 0 && row < r_num && col >= 0 && col < c_num) {
|
||||
PVector accu = motion_field.get(row * c_num + col);
|
||||
accu.x += p1.x - p0.x;
|
||||
accu.y += p1.y - p0.y;
|
||||
accu.z += 1;
|
||||
}
|
||||
}
|
||||
// if some blocks do not have point, then use ray tracing to see if they are
|
||||
// in triangles
|
||||
for (int i = 0; i < r_num; i++)
|
||||
for (int j = 0; j < c_num; j++) {
|
||||
PVector accu = motion_field.get(i * c_num + j);
|
||||
if (accu.z > 0) {
|
||||
continue;
|
||||
}
|
||||
// use the center of the block to generate view ray
|
||||
float cx = j * block_size + block_size / 2.0f - width / 2.0f;
|
||||
float cy = i * block_size + block_size / 2.0f - height / 2.0f;
|
||||
float cz = 0.5f * height / tan(current_cam.fov / 2.0f);
|
||||
PVector dir = new PVector(cx, cy, cz);
|
||||
float[] camMat = current_cam.getCameraMat();
|
||||
dir = MatxVec3(transpose3x3(camMat), dir);
|
||||
dir.normalize();
|
||||
Ray r = new Ray(current_cam.pos, dir);
|
||||
// ray tracing
|
||||
float[] param = new float[4];
|
||||
param[0] = Float.POSITIVE_INFINITY;
|
||||
if (bvh.intersect(r, param)) {
|
||||
PVector p = new PVector(param[1], param[2], param[3]);
|
||||
PVector p0 = current_cam.project(p);
|
||||
PVector p1 = last_cam.project(p);
|
||||
accu.x += p1.x - p0.x;
|
||||
accu.y += p1.y - p0.y;
|
||||
accu.z += 1;
|
||||
}
|
||||
}
|
||||
// estimate the motion vector of each block
|
||||
for (int i = 0; i < r_num * c_num; i++) {
|
||||
PVector mv = motion_field.get(i);
|
||||
if (mv.z > 0) {
|
||||
motion_field.set(i, new PVector(mv.x / mv.z, mv.y / mv.z, 0));
|
||||
} else // there is nothing in the block, use -1 to mark it.
|
||||
{
|
||||
motion_field.set(i, new PVector(0.0, 0.0, -1));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void render() {
|
||||
int r_num = height / block_size, c_num = width / block_size;
|
||||
for (int i = 0; i < r_num; i++)
|
||||
for (int j = 0; j < c_num; j++) {
|
||||
PVector mv = motion_field.get(i * c_num + j);
|
||||
float ox = j * block_size + 0.5f * block_size;
|
||||
float oy = i * block_size + 0.5f * block_size;
|
||||
stroke(255, 0, 0);
|
||||
line(ox, oy, ox + mv.x, oy + mv.y);
|
||||
}
|
||||
}
|
||||
|
||||
void save(String path) {
|
||||
int r_num = height / block_size;
|
||||
int c_num = width / block_size;
|
||||
String[] mvs = new String[r_num];
|
||||
for (int i = 0; i < r_num; i++) {
|
||||
mvs[i] = "";
|
||||
for (int j = 0; j < c_num; j++) {
|
||||
PVector mv = motion_field.get(i * c_num + j);
|
||||
if (mv.z != -1) {
|
||||
mvs[i] += str(mv.x) + "," + str(mv.y);
|
||||
} else // there is nothing
|
||||
{
|
||||
mvs[i] += "-,-";
|
||||
}
|
||||
if (j != c_num - 1) mvs[i] += ";";
|
||||
}
|
||||
}
|
||||
saveStrings(path, mvs);
|
||||
}
|
||||
}
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
class PointCloud {
|
||||
ArrayList<PVector> points; // array to save points
|
||||
IntList point_colors; // array to save points color
|
||||
PVector cloud_mass;
|
||||
float[] depth;
|
||||
boolean[] real;
|
||||
PointCloud() {
|
||||
// initialize
|
||||
points = new ArrayList<PVector>();
|
||||
point_colors = new IntList();
|
||||
cloud_mass = new PVector(0, 0, 0);
|
||||
depth = new float[width * height];
|
||||
real = new boolean[width * height];
|
||||
}
|
||||
|
||||
void generate(PImage rgb_img, PImage depth_img, Transform trans) {
|
||||
if (depth_img.width != width || depth_img.height != height ||
|
||||
rgb_img.width != width || rgb_img.height != height) {
|
||||
println("rgb and depth file dimension should be same with window size");
|
||||
exit();
|
||||
}
|
||||
// clear depth and real
|
||||
for (int i = 0; i < width * height; i++) {
|
||||
depth[i] = 0;
|
||||
real[i] = false;
|
||||
}
|
||||
for (int v = 0; v < height; v++)
|
||||
for (int u = 0; u < width; u++) {
|
||||
// get depth value (red channel)
|
||||
color depth_px = depth_img.get(u, v);
|
||||
depth[v * width + u] = depth_px & 0x0000FFFF;
|
||||
if (int(depth[v * width + u]) != 0) {
|
||||
real[v * width + u] = true;
|
||||
}
|
||||
point_colors.append(rgb_img.get(u, v));
|
||||
}
|
||||
for (int v = 0; v < height; v++)
|
||||
for (int u = 0; u < width; u++) {
|
||||
if (int(depth[v * width + u]) == 0) {
|
||||
interpolateDepth(v, u);
|
||||
}
|
||||
// add transformed pixel as well as pixel color to the list
|
||||
PVector pos = trans.transform(u, v, int(depth[v * width + u]));
|
||||
points.add(pos);
|
||||
// accumulate z value
|
||||
cloud_mass = PVector.add(cloud_mass, pos);
|
||||
}
|
||||
}
|
||||
void fillInDepthAlongPath(float d, Node node) {
|
||||
node = node.parent;
|
||||
while (node != null) {
|
||||
int i = node.row;
|
||||
int j = node.col;
|
||||
if (depth[i * width + j] == 0) {
|
||||
depth[i * width + j] = d;
|
||||
}
|
||||
node = node.parent;
|
||||
}
|
||||
}
|
||||
// interpolate
|
||||
void interpolateDepth(int row, int col) {
|
||||
if (row < 0 || row >= height || col < 0 || col >= width ||
|
||||
int(depth[row * width + col]) != 0) {
|
||||
return;
|
||||
}
|
||||
ArrayList<Node> queue = new ArrayList<Node>();
|
||||
queue.add(new Node(row, col, null));
|
||||
boolean[] visited = new boolean[width * height];
|
||||
for (int i = 0; i < width * height; i++) visited[i] = false;
|
||||
visited[row * width + col] = true;
|
||||
// Using BFS to Find the Nearest Neighbor
|
||||
while (queue.size() > 0) {
|
||||
// pop
|
||||
Node node = queue.get(0);
|
||||
queue.remove(0);
|
||||
int i = node.row;
|
||||
int j = node.col;
|
||||
// if current position have a real depth
|
||||
if (depth[i * width + j] != 0 && real[i * width + j]) {
|
||||
fillInDepthAlongPath(depth[i * width + j], node);
|
||||
break;
|
||||
} else {
|
||||
// search unvisited 8 neighbors
|
||||
for (int r = max(0, i - 1); r < min(height, i + 2); r++) {
|
||||
for (int c = max(0, j - 1); c < min(width, j + 2); c++) {
|
||||
if (!visited[r * width + c]) {
|
||||
visited[r * width + c] = true;
|
||||
queue.add(new Node(r, c, node));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// get point cloud size
|
||||
int size() { return points.size(); }
|
||||
// get ith position
|
||||
PVector getPosition(int i) {
|
||||
if (i >= points.size()) {
|
||||
println("point position: index " + str(i) + " exceeds");
|
||||
exit();
|
||||
}
|
||||
return points.get(i);
|
||||
}
|
||||
// get ith color
|
||||
color getColor(int i) {
|
||||
if (i >= point_colors.size()) {
|
||||
println("point color: index " + str(i) + " exceeds");
|
||||
exit();
|
||||
}
|
||||
return point_colors.get(i);
|
||||
}
|
||||
// get cloud center
|
||||
PVector getCloudCenter() {
|
||||
if (points.size() > 0) {
|
||||
return PVector.div(cloud_mass, points.size());
|
||||
}
|
||||
return new PVector(0, 0, 0);
|
||||
}
|
||||
// merge two clouds
|
||||
void merge(PointCloud point_cloud) {
|
||||
for (int i = 0; i < point_cloud.size(); i++) {
|
||||
points.add(point_cloud.getPosition(i));
|
||||
point_colors.append(point_cloud.getColor(i));
|
||||
}
|
||||
cloud_mass = PVector.add(cloud_mass, point_cloud.cloud_mass);
|
||||
}
|
||||
}
|
||||
|
||||
class Node {
|
||||
int row, col;
|
||||
Node parent;
|
||||
Node(int row, int col, Node parent) {
|
||||
this.row = row;
|
||||
this.col = col;
|
||||
this.parent = parent;
|
||||
}
|
||||
}
|
||||
+61
@@ -0,0 +1,61 @@
|
||||
// Triangle
|
||||
class Triangle {
|
||||
// position
|
||||
PVector p1, p2, p3;
|
||||
// color
|
||||
color c1, c2, c3;
|
||||
BoundingBox bbx;
|
||||
Triangle(PVector p1, PVector p2, PVector p3, color c1, color c2, color c3) {
|
||||
this.p1 = p1;
|
||||
this.p2 = p2;
|
||||
this.p3 = p3;
|
||||
this.c1 = c1;
|
||||
this.c2 = c2;
|
||||
this.c3 = c3;
|
||||
bbx = new BoundingBox();
|
||||
bbx.create(this);
|
||||
}
|
||||
// check to see if a ray intersects with the triangle
|
||||
boolean intersect(Ray r, float[] param) {
|
||||
PVector p21 = PVector.sub(p2, p1);
|
||||
PVector p31 = PVector.sub(p3, p1);
|
||||
PVector po1 = PVector.sub(r.ori, p1);
|
||||
|
||||
PVector dxp31 = r.dir.cross(p31);
|
||||
PVector po1xp21 = po1.cross(p21);
|
||||
float denom = p21.dot(dxp31);
|
||||
float t = p31.dot(po1xp21) / denom;
|
||||
float alpha = po1.dot(dxp31) / denom;
|
||||
float beta = r.dir.dot(po1xp21) / denom;
|
||||
|
||||
boolean res = t > 0 && alpha > 0 && alpha < 1 && beta > 0 && beta < 1 &&
|
||||
alpha + beta < 1;
|
||||
// depth test
|
||||
if (res && t < param[0]) {
|
||||
param[0] = t;
|
||||
param[1] = alpha * p1.x + beta * p2.x + (1 - alpha - beta) * p3.x;
|
||||
param[2] = alpha * p1.y + beta * p2.y + (1 - alpha - beta) * p3.y;
|
||||
param[3] = alpha * p1.z + beta * p2.z + (1 - alpha - beta) * p3.z;
|
||||
}
|
||||
return res;
|
||||
}
|
||||
void render() {
|
||||
beginShape(TRIANGLES);
|
||||
fill(c1);
|
||||
vertex(p1.x, p1.y, p1.z);
|
||||
fill(c2);
|
||||
vertex(p2.x, p2.y, p2.z);
|
||||
fill(c3);
|
||||
vertex(p3.x, p3.y, p3.z);
|
||||
endShape();
|
||||
}
|
||||
}
|
||||
// Ray
|
||||
class Ray {
|
||||
// origin and direction
|
||||
PVector ori, dir;
|
||||
Ray(PVector ori, PVector dir) {
|
||||
this.ori = ori;
|
||||
this.dir = dir;
|
||||
}
|
||||
}
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
class Scene {
|
||||
PointCloud point_cloud;
|
||||
ArrayList<Triangle> mesh;
|
||||
BVH bvh;
|
||||
MotionField motion_field;
|
||||
Camera last_cam;
|
||||
Camera current_cam;
|
||||
int frame_count;
|
||||
|
||||
Scene(Camera camera, PointCloud point_cloud, MotionField motion_field) {
|
||||
this.point_cloud = point_cloud;
|
||||
this.motion_field = motion_field;
|
||||
mesh = new ArrayList<Triangle>();
|
||||
for (int v = 0; v < height - 1; v++)
|
||||
for (int u = 0; u < width - 1; u++) {
|
||||
PVector p1 = point_cloud.getPosition(v * width + u);
|
||||
PVector p2 = point_cloud.getPosition(v * width + u + 1);
|
||||
PVector p3 = point_cloud.getPosition((v + 1) * width + u + 1);
|
||||
PVector p4 = point_cloud.getPosition((v + 1) * width + u);
|
||||
color c1 = point_cloud.getColor(v * width + u);
|
||||
color c2 = point_cloud.getColor(v * width + u + 1);
|
||||
color c3 = point_cloud.getColor((v + 1) * width + u + 1);
|
||||
color c4 = point_cloud.getColor((v + 1) * width + u);
|
||||
mesh.add(new Triangle(p1, p2, p3, c1, c2, c3));
|
||||
mesh.add(new Triangle(p3, p4, p1, c3, c4, c1));
|
||||
}
|
||||
bvh = new BVH(mesh);
|
||||
last_cam = camera.copy();
|
||||
current_cam = camera;
|
||||
frame_count = 0;
|
||||
}
|
||||
|
||||
void run() {
|
||||
last_cam = current_cam.copy();
|
||||
current_cam.run();
|
||||
motion_field.update(last_cam, current_cam, point_cloud, bvh);
|
||||
frame_count += 1;
|
||||
}
|
||||
|
||||
void render(boolean show_motion_field) {
|
||||
// build mesh
|
||||
current_cam.open();
|
||||
noStroke();
|
||||
for (int i = 0; i < mesh.size(); i++) {
|
||||
Triangle t = mesh.get(i);
|
||||
t.render();
|
||||
}
|
||||
if (show_motion_field) {
|
||||
current_cam.close();
|
||||
motion_field.render();
|
||||
}
|
||||
}
|
||||
|
||||
void save(String path) { saveFrame(path + "_" + str(frame_count) + ".png"); }
|
||||
|
||||
void saveMotionField(String path) {
|
||||
motion_field.save(path + "_" + str(frame_count) + ".txt");
|
||||
}
|
||||
}
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
class Transform {
|
||||
float[] inv_rot; // inverse of rotation matrix
|
||||
PVector inv_mov; // inverse of movement vector
|
||||
float focal; // the focal distacne of real camera
|
||||
int w, h; // the width and height of the frame
|
||||
float normalier; // nomalization factor of depth
|
||||
Transform(float tx, float ty, float tz, float qx, float qy, float qz,
|
||||
float qw, float fov, int w, int h, float normalier) {
|
||||
// currently, we did not use the info of real camera's position and
|
||||
// quaternion maybe we will use it in the future when combine all frames
|
||||
float[] rot = quaternion2Mat3x3(qx, qy, qz, qw);
|
||||
inv_rot = transpose3x3(rot);
|
||||
inv_mov = new PVector(-tx, -ty, -tz);
|
||||
this.focal = 0.5f * h / tan(fov / 2.0);
|
||||
this.w = w;
|
||||
this.h = h;
|
||||
this.normalier = normalier;
|
||||
}
|
||||
|
||||
PVector transform(int i, int j, float d) {
|
||||
// transfer from camera view to world view
|
||||
float z = d / normalier;
|
||||
float x = (i - w / 2.0f) * z / focal;
|
||||
float y = (j - h / 2.0f) * z / focal;
|
||||
return new PVector(x, y, z);
|
||||
}
|
||||
}
|
||||
|
||||
// get rotation matrix by using rotation axis and angle
|
||||
float[] getRotationMat3x3(float angle, float ax, float ay, float az) {
|
||||
float[] mat = new float[9];
|
||||
float c = cos(angle);
|
||||
float s = sin(angle);
|
||||
mat[0] = c + ax * ax * (1 - c);
|
||||
mat[1] = ax * ay * (1 - c) - az * s;
|
||||
mat[2] = ax * az * (1 - c) + ay * s;
|
||||
mat[3] = ay * ax * (1 - c) + az * s;
|
||||
mat[4] = c + ay * ay * (1 - c);
|
||||
mat[5] = ay * az * (1 - c) - ax * s;
|
||||
mat[6] = az * ax * (1 - c) - ay * s;
|
||||
mat[7] = az * ay * (1 - c) + ax * s;
|
||||
mat[8] = c + az * az * (1 - c);
|
||||
return mat;
|
||||
}
|
||||
|
||||
// get rotation matrix by using quaternion
|
||||
float[] quaternion2Mat3x3(float qx, float qy, float qz, float qw) {
|
||||
float[] mat = new float[9];
|
||||
mat[0] = 1 - 2 * qy * qy - 2 * qz * qz;
|
||||
mat[1] = 2 * qx * qy - 2 * qz * qw;
|
||||
mat[2] = 2 * qx * qz + 2 * qy * qw;
|
||||
mat[3] = 2 * qx * qy + 2 * qz * qw;
|
||||
mat[4] = 1 - 2 * qx * qx - 2 * qz * qz;
|
||||
mat[5] = 2 * qy * qz - 2 * qx * qw;
|
||||
mat[6] = 2 * qx * qz - 2 * qy * qw;
|
||||
mat[7] = 2 * qy * qz + 2 * qx * qw;
|
||||
mat[8] = 1 - 2 * qx * qx - 2 * qy * qy;
|
||||
return mat;
|
||||
}
|
||||
|
||||
// tranpose a 3x3 matrix
|
||||
float[] transpose3x3(float[] mat) {
|
||||
float[] Tmat = new float[9];
|
||||
for (int i = 0; i < 3; i++)
|
||||
for (int j = 0; j < 3; j++) {
|
||||
Tmat[i * 3 + j] = mat[j * 3 + i];
|
||||
}
|
||||
return Tmat;
|
||||
}
|
||||
|
||||
// multiply a matrix with vector
|
||||
PVector MatxVec3(float[] mat, PVector v) {
|
||||
float[] vec = v.array();
|
||||
float[] res = new float[3];
|
||||
for (int i = 0; i < 3; i++) {
|
||||
res[i] = 0.0f;
|
||||
for (int j = 0; j < 3; j++) {
|
||||
res[i] += mat[i * 3 + j] * vec[j];
|
||||
}
|
||||
}
|
||||
return new PVector(res[0], res[1], res[2]);
|
||||
}
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
// show grids
|
||||
void showGrids(int block_size) {
|
||||
ortho(-width, 0, -height, 0);
|
||||
camera(0, 0, 0, 0, 0, 1, 0, 1, 0);
|
||||
stroke(0, 0, 255);
|
||||
for (int i = 0; i < height; i += block_size) {
|
||||
line(0, i, width, i);
|
||||
}
|
||||
for (int i = 0; i < width; i += block_size) {
|
||||
line(i, 0, i, height);
|
||||
}
|
||||
}
|
||||
|
||||
// save the point clould information
|
||||
void savePointCloud(PointCloud point_cloud, String file_name) {
|
||||
String[] positions = new String[point_cloud.points.size()];
|
||||
String[] colors = new String[point_cloud.points.size()];
|
||||
for (int i = 0; i < point_cloud.points.size(); i++) {
|
||||
PVector point = point_cloud.getPosition(i);
|
||||
color point_color = point_cloud.getColor(i);
|
||||
positions[i] = str(point.x) + ' ' + str(point.y) + ' ' + str(point.z);
|
||||
colors[i] = str(((point_color >> 16) & 0xFF) / 255.0) + ' ' +
|
||||
str(((point_color >> 8) & 0xFF) / 255.0) + ' ' +
|
||||
str((point_color & 0xFF) / 255.0);
|
||||
}
|
||||
saveStrings(file_name + "_pos.txt", positions);
|
||||
saveStrings(file_name + "_color.txt", colors);
|
||||
}
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
/*The dataset is from
|
||||
*Computer Vision Group
|
||||
*TUM Department of Informatics Technical
|
||||
*University of Munich
|
||||
*https://vision.in.tum.de/data/datasets/rgbd-dataset/download#freiburg1_xyz
|
||||
*/
|
||||
Scene scene;
|
||||
void setup() {
|
||||
size(640, 480, P3D);
|
||||
// default settings
|
||||
int frame_no = 0; // frame number
|
||||
float fov = PI / 3; // field of view
|
||||
int block_size = 8; // block size
|
||||
float normalizer = 5000.0f; // normalizer
|
||||
// initialize
|
||||
PointCloud point_cloud = new PointCloud();
|
||||
// synchronized rgb, depth and ground truth
|
||||
String head = "../data/";
|
||||
String[] rgb_depth_gt = loadStrings(head + "rgb_depth_groundtruth.txt");
|
||||
// read in rgb and depth image file paths as well as corresponding camera
|
||||
// posiiton and quaternion
|
||||
String[] info = split(rgb_depth_gt[frame_no], ' ');
|
||||
String rgb_path = head + info[1];
|
||||
String depth_path = head + info[3];
|
||||
float tx = float(info[7]), ty = float(info[8]),
|
||||
tz = float(info[9]); // real camera position
|
||||
float qx = float(info[10]), qy = float(info[11]), qz = float(info[12]),
|
||||
qw = float(info[13]); // quaternion
|
||||
|
||||
// build transformer
|
||||
Transform trans =
|
||||
new Transform(tx, ty, tz, qx, qy, qz, qw, fov, width, height, normalizer);
|
||||
PImage rgb = loadImage(rgb_path);
|
||||
PImage depth = loadImage(depth_path);
|
||||
// generate point cloud
|
||||
point_cloud.generate(rgb, depth, trans);
|
||||
// initialize camera
|
||||
Camera camera = new Camera(fov, new PVector(0, 0, 0), new PVector(0, 0, 1),
|
||||
new PVector(0, 1, 0));
|
||||
// initialize motion field
|
||||
MotionField motion_field = new MotionField(block_size);
|
||||
// initialize scene
|
||||
scene = new Scene(camera, point_cloud, motion_field);
|
||||
}
|
||||
boolean inter = false;
|
||||
void draw() {
|
||||
background(0);
|
||||
// run camera dragged mouse to rotate camera
|
||||
// w: go forward
|
||||
// s: go backward
|
||||
// a: go left
|
||||
// d: go right
|
||||
// up arrow: go up
|
||||
// down arrow: go down
|
||||
//+ increase move speed
|
||||
//- decrease move speed
|
||||
// r: rotate the camera
|
||||
// b: reset to initial position
|
||||
scene.run(); // true: make interpolation; false: do not make
|
||||
// interpolation
|
||||
if (keyPressed && key == 'o') {
|
||||
inter = true;
|
||||
}
|
||||
scene.render(
|
||||
false); // true: turn on motion field; false: turn off motion field
|
||||
// save frame with no motion field
|
||||
scene.save("../data/frame/raw");
|
||||
background(0);
|
||||
scene.render(true);
|
||||
showGrids(scene.motion_field.block_size);
|
||||
// save frame with motion field
|
||||
scene.save("../data/frame/raw_mv");
|
||||
scene.saveMotionField("../data/frame/mv");
|
||||
}
|
||||
+4756
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,130 @@
|
||||
#!/usr/bin/env python
|
||||
## Copyright (c) 2012 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.
|
||||
##
|
||||
"""Classes for representing diff pieces."""
|
||||
|
||||
__author__ = "jkoleszar@google.com"
|
||||
|
||||
import re
|
||||
|
||||
|
||||
class DiffLines(object):
|
||||
"""A container for one half of a diff."""
|
||||
|
||||
def __init__(self, filename, offset, length):
|
||||
self.filename = filename
|
||||
self.offset = offset
|
||||
self.length = length
|
||||
self.lines = []
|
||||
self.delta_line_nums = []
|
||||
|
||||
def Append(self, line):
|
||||
l = len(self.lines)
|
||||
if line[0] != " ":
|
||||
self.delta_line_nums.append(self.offset + l)
|
||||
self.lines.append(line[1:])
|
||||
assert l+1 <= self.length
|
||||
|
||||
def Complete(self):
|
||||
return len(self.lines) == self.length
|
||||
|
||||
def __contains__(self, item):
|
||||
return item >= self.offset and item <= self.offset + self.length - 1
|
||||
|
||||
|
||||
class DiffHunk(object):
|
||||
"""A container for one diff hunk, consisting of two DiffLines."""
|
||||
|
||||
def __init__(self, header, file_a, file_b, start_a, len_a, start_b, len_b):
|
||||
self.header = header
|
||||
self.left = DiffLines(file_a, start_a, len_a)
|
||||
self.right = DiffLines(file_b, start_b, len_b)
|
||||
self.lines = []
|
||||
|
||||
def Append(self, line):
|
||||
"""Adds a line to the DiffHunk and its DiffLines children."""
|
||||
if line[0] == "-":
|
||||
self.left.Append(line)
|
||||
elif line[0] == "+":
|
||||
self.right.Append(line)
|
||||
elif line[0] == " ":
|
||||
self.left.Append(line)
|
||||
self.right.Append(line)
|
||||
elif line[0] == "\\":
|
||||
# Ignore newline messages from git diff.
|
||||
pass
|
||||
else:
|
||||
assert False, ("Unrecognized character at start of diff line "
|
||||
"%r" % line[0])
|
||||
self.lines.append(line)
|
||||
|
||||
def Complete(self):
|
||||
return self.left.Complete() and self.right.Complete()
|
||||
|
||||
def __repr__(self):
|
||||
return "DiffHunk(%s, %s, len %d)" % (
|
||||
self.left.filename, self.right.filename,
|
||||
max(self.left.length, self.right.length))
|
||||
|
||||
|
||||
def ParseDiffHunks(stream):
|
||||
"""Walk a file-like object, yielding DiffHunks as they're parsed."""
|
||||
|
||||
file_regex = re.compile(r"(\+\+\+|---) (\S+)")
|
||||
range_regex = re.compile(r"@@ -(\d+)(,(\d+))? \+(\d+)(,(\d+))?")
|
||||
hunk = None
|
||||
while True:
|
||||
line = stream.readline()
|
||||
if not line:
|
||||
break
|
||||
|
||||
if hunk is None:
|
||||
# Parse file names
|
||||
diff_file = file_regex.match(line)
|
||||
if diff_file:
|
||||
if line.startswith("---"):
|
||||
a_line = line
|
||||
a = diff_file.group(2)
|
||||
continue
|
||||
if line.startswith("+++"):
|
||||
b_line = line
|
||||
b = diff_file.group(2)
|
||||
continue
|
||||
|
||||
# Parse offset/lengths
|
||||
diffrange = range_regex.match(line)
|
||||
if diffrange:
|
||||
if diffrange.group(2):
|
||||
start_a = int(diffrange.group(1))
|
||||
len_a = int(diffrange.group(3))
|
||||
else:
|
||||
start_a = 1
|
||||
len_a = int(diffrange.group(1))
|
||||
|
||||
if diffrange.group(5):
|
||||
start_b = int(diffrange.group(4))
|
||||
len_b = int(diffrange.group(6))
|
||||
else:
|
||||
start_b = 1
|
||||
len_b = int(diffrange.group(4))
|
||||
|
||||
header = [a_line, b_line, line]
|
||||
hunk = DiffHunk(header, a, b, start_a, len_a, start_b, len_b)
|
||||
else:
|
||||
# Add the current line to the hunk
|
||||
hunk.Append(line)
|
||||
|
||||
# See if the whole hunk has been parsed. If so, yield it and prepare
|
||||
# for the next hunk.
|
||||
if hunk.Complete():
|
||||
yield hunk
|
||||
hunk = None
|
||||
|
||||
# Partial hunks are a parse error
|
||||
assert hunk is None
|
||||
@@ -0,0 +1,14 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Add organization names manually.
|
||||
|
||||
cat <<EOF
|
||||
# This file is automatically generated from the git commit history
|
||||
# by tools/gen_authors.sh.
|
||||
|
||||
$(git log --pretty=format:"%aN <%aE>" | sort | uniq | grep -v corp.google \
|
||||
| grep -v noreply)
|
||||
Google Inc.
|
||||
The Mozilla Foundation
|
||||
The Xiph.Org Foundation
|
||||
EOF
|
||||
@@ -0,0 +1,76 @@
|
||||
#!/usr/bin/env python
|
||||
## Copyright (c) 2012 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.
|
||||
##
|
||||
"""Calculates the "intersection" of two unified diffs.
|
||||
|
||||
Given two diffs, A and B, it finds all hunks in B that had non-context lines
|
||||
in A and prints them to stdout. This is useful to determine the hunks in B that
|
||||
are relevant to A. The resulting file can be applied with patch(1) on top of A.
|
||||
"""
|
||||
|
||||
__author__ = "jkoleszar@google.com"
|
||||
|
||||
import sys
|
||||
|
||||
import diff
|
||||
|
||||
|
||||
def FormatDiffHunks(hunks):
|
||||
"""Re-serialize a list of DiffHunks."""
|
||||
r = []
|
||||
last_header = None
|
||||
for hunk in hunks:
|
||||
this_header = hunk.header[0:2]
|
||||
if last_header != this_header:
|
||||
r.extend(hunk.header)
|
||||
last_header = this_header
|
||||
else:
|
||||
r.extend(hunk.header[2])
|
||||
r.extend(hunk.lines)
|
||||
r.append("\n")
|
||||
return "".join(r)
|
||||
|
||||
|
||||
def ZipHunks(rhs_hunks, lhs_hunks):
|
||||
"""Join two hunk lists on filename."""
|
||||
for rhs_hunk in rhs_hunks:
|
||||
rhs_file = rhs_hunk.right.filename.split("/")[1:]
|
||||
|
||||
for lhs_hunk in lhs_hunks:
|
||||
lhs_file = lhs_hunk.left.filename.split("/")[1:]
|
||||
if lhs_file != rhs_file:
|
||||
continue
|
||||
yield (rhs_hunk, lhs_hunk)
|
||||
|
||||
|
||||
def main():
|
||||
old_hunks = [x for x in diff.ParseDiffHunks(open(sys.argv[1], "r"))]
|
||||
new_hunks = [x for x in diff.ParseDiffHunks(open(sys.argv[2], "r"))]
|
||||
out_hunks = []
|
||||
|
||||
# Join the right hand side of the older diff with the left hand side of the
|
||||
# newer diff.
|
||||
for old_hunk, new_hunk in ZipHunks(old_hunks, new_hunks):
|
||||
if new_hunk in out_hunks:
|
||||
continue
|
||||
old_lines = old_hunk.right
|
||||
new_lines = new_hunk.left
|
||||
|
||||
# Determine if this hunk overlaps any non-context line from the other
|
||||
for i in old_lines.delta_line_nums:
|
||||
if i in new_lines:
|
||||
out_hunks.append(new_hunk)
|
||||
break
|
||||
|
||||
if out_hunks:
|
||||
print FormatDiffHunks(out_hunks)
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,144 @@
|
||||
#!/usr/bin/python
|
||||
## Copyright (c) 2012 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.
|
||||
##
|
||||
"""Performs style checking on each diff hunk."""
|
||||
import getopt
|
||||
import os
|
||||
import StringIO
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import diff
|
||||
|
||||
|
||||
SHORT_OPTIONS = "h"
|
||||
LONG_OPTIONS = ["help"]
|
||||
|
||||
TOPLEVEL_CMD = ["git", "rev-parse", "--show-toplevel"]
|
||||
DIFF_CMD = ["git", "diff"]
|
||||
DIFF_INDEX_CMD = ["git", "diff-index", "-u", "HEAD", "--"]
|
||||
SHOW_CMD = ["git", "show"]
|
||||
CPPLINT_FILTERS = ["-readability/casting"]
|
||||
|
||||
|
||||
class Usage(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class SubprocessException(Exception):
|
||||
def __init__(self, args):
|
||||
msg = "Failed to execute '%s'"%(" ".join(args))
|
||||
super(SubprocessException, self).__init__(msg)
|
||||
|
||||
|
||||
class Subprocess(subprocess.Popen):
|
||||
"""Adds the notion of an expected returncode to Popen."""
|
||||
|
||||
def __init__(self, args, expected_returncode=0, **kwargs):
|
||||
self._args = args
|
||||
self._expected_returncode = expected_returncode
|
||||
super(Subprocess, self).__init__(args, **kwargs)
|
||||
|
||||
def communicate(self, *args, **kwargs):
|
||||
result = super(Subprocess, self).communicate(*args, **kwargs)
|
||||
if self._expected_returncode is not None:
|
||||
try:
|
||||
ok = self.returncode in self._expected_returncode
|
||||
except TypeError:
|
||||
ok = self.returncode == self._expected_returncode
|
||||
if not ok:
|
||||
raise SubprocessException(self._args)
|
||||
return result
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
if argv is None:
|
||||
argv = sys.argv
|
||||
try:
|
||||
try:
|
||||
opts, args = getopt.getopt(argv[1:], SHORT_OPTIONS, LONG_OPTIONS)
|
||||
except getopt.error, msg:
|
||||
raise Usage(msg)
|
||||
|
||||
# process options
|
||||
for o, _ in opts:
|
||||
if o in ("-h", "--help"):
|
||||
print __doc__
|
||||
sys.exit(0)
|
||||
|
||||
if args and len(args) > 1:
|
||||
print __doc__
|
||||
sys.exit(0)
|
||||
|
||||
# Find the fully qualified path to the root of the tree
|
||||
tl = Subprocess(TOPLEVEL_CMD, stdout=subprocess.PIPE)
|
||||
tl = tl.communicate()[0].strip()
|
||||
|
||||
# See if we're working on the index or not.
|
||||
if args:
|
||||
diff_cmd = DIFF_CMD + [args[0] + "^!"]
|
||||
else:
|
||||
diff_cmd = DIFF_INDEX_CMD
|
||||
|
||||
# Build the command line to execute cpplint
|
||||
cpplint_cmd = [os.path.join(tl, "tools", "cpplint.py"),
|
||||
"--filter=" + ",".join(CPPLINT_FILTERS),
|
||||
"-"]
|
||||
|
||||
# Get a list of all affected lines
|
||||
file_affected_line_map = {}
|
||||
p = Subprocess(diff_cmd, stdout=subprocess.PIPE)
|
||||
stdout = p.communicate()[0]
|
||||
for hunk in diff.ParseDiffHunks(StringIO.StringIO(stdout)):
|
||||
filename = hunk.right.filename[2:]
|
||||
if filename not in file_affected_line_map:
|
||||
file_affected_line_map[filename] = set()
|
||||
file_affected_line_map[filename].update(hunk.right.delta_line_nums)
|
||||
|
||||
# Run each affected file through cpplint
|
||||
lint_failed = False
|
||||
for filename, affected_lines in file_affected_line_map.iteritems():
|
||||
if filename.split(".")[-1] not in ("c", "h", "cc"):
|
||||
continue
|
||||
|
||||
if args:
|
||||
# File contents come from git
|
||||
show_cmd = SHOW_CMD + [args[0] + ":" + filename]
|
||||
show = Subprocess(show_cmd, stdout=subprocess.PIPE)
|
||||
lint = Subprocess(cpplint_cmd, expected_returncode=(0, 1),
|
||||
stdin=show.stdout, stderr=subprocess.PIPE)
|
||||
lint_out = lint.communicate()[1]
|
||||
else:
|
||||
# File contents come from the working tree
|
||||
lint = Subprocess(cpplint_cmd, expected_returncode=(0, 1),
|
||||
stdin=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
stdin = open(os.path.join(tl, filename)).read()
|
||||
lint_out = lint.communicate(stdin)[1]
|
||||
|
||||
for line in lint_out.split("\n"):
|
||||
fields = line.split(":")
|
||||
if fields[0] != "-":
|
||||
continue
|
||||
warning_line_num = int(fields[1])
|
||||
if warning_line_num in affected_lines:
|
||||
print "%s:%d:%s"%(filename, warning_line_num,
|
||||
":".join(fields[2:]))
|
||||
lint_failed = True
|
||||
|
||||
# Set exit code if any relevant lint errors seen
|
||||
if lint_failed:
|
||||
return 1
|
||||
|
||||
except Usage, err:
|
||||
print >>sys.stderr, err
|
||||
print >>sys.stderr, "for help use --help"
|
||||
return 2
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,186 @@
|
||||
import sys
|
||||
import matplotlib.pyplot as plt
|
||||
from matplotlib.collections import LineCollection
|
||||
from matplotlib import colors as mcolors
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
|
||||
def draw_mv_ls(axis, mv_ls, mode=0):
|
||||
colors = np.array([(1., 0., 0., 1.)])
|
||||
segs = np.array([
|
||||
np.array([[ptr[0], ptr[1]], [ptr[0] + ptr[2], ptr[1] + ptr[3]]])
|
||||
for ptr in mv_ls
|
||||
])
|
||||
line_segments = LineCollection(
|
||||
segs, linewidths=(1.,), colors=colors, linestyle='solid')
|
||||
axis.add_collection(line_segments)
|
||||
if mode == 0:
|
||||
axis.scatter(mv_ls[:, 0], mv_ls[:, 1], s=2, c='b')
|
||||
else:
|
||||
axis.scatter(
|
||||
mv_ls[:, 0] + mv_ls[:, 2], mv_ls[:, 1] + mv_ls[:, 3], s=2, c='b')
|
||||
|
||||
|
||||
def draw_pred_block_ls(axis, mv_ls, bs, mode=0):
|
||||
colors = np.array([(0., 0., 0., 1.)])
|
||||
segs = []
|
||||
for ptr in mv_ls:
|
||||
if mode == 0:
|
||||
x = ptr[0]
|
||||
y = ptr[1]
|
||||
else:
|
||||
x = ptr[0] + ptr[2]
|
||||
y = ptr[1] + ptr[3]
|
||||
x_ls = [x, x + bs, x + bs, x, x]
|
||||
y_ls = [y, y, y + bs, y + bs, y]
|
||||
|
||||
segs.append(np.column_stack([x_ls, y_ls]))
|
||||
line_segments = LineCollection(
|
||||
segs, linewidths=(.5,), colors=colors, linestyle='solid')
|
||||
axis.add_collection(line_segments)
|
||||
|
||||
|
||||
def read_frame(fp, no_swap=0):
|
||||
plane = [None, None, None]
|
||||
for i in range(3):
|
||||
line = fp.readline()
|
||||
word_ls = line.split()
|
||||
word_ls = [int(item) for item in word_ls]
|
||||
rows = word_ls[0]
|
||||
cols = word_ls[1]
|
||||
|
||||
line = fp.readline()
|
||||
word_ls = line.split()
|
||||
word_ls = [int(item) for item in word_ls]
|
||||
|
||||
plane[i] = np.array(word_ls).reshape(rows, cols)
|
||||
if i > 0:
|
||||
plane[i] = plane[i].repeat(2, axis=0).repeat(2, axis=1)
|
||||
plane = np.array(plane)
|
||||
if no_swap == 0:
|
||||
plane = np.swapaxes(np.swapaxes(plane, 0, 1), 1, 2)
|
||||
return plane
|
||||
|
||||
|
||||
def yuv_to_rgb(yuv):
|
||||
#mat = np.array([
|
||||
# [1.164, 0 , 1.596 ],
|
||||
# [1.164, -0.391, -0.813],
|
||||
# [1.164, 2.018 , 0 ] ]
|
||||
# )
|
||||
#c = np.array([[ -16 , -16 , -16 ],
|
||||
# [ 0 , -128, -128 ],
|
||||
# [ -128, -128, 0 ]])
|
||||
|
||||
mat = np.array([[1, 0, 1.4075], [1, -0.3445, -0.7169], [1, 1.7790, 0]])
|
||||
c = np.array([[0, 0, 0], [0, -128, -128], [-128, -128, 0]])
|
||||
mat_c = np.dot(mat, c)
|
||||
v = np.array([mat_c[0, 0], mat_c[1, 1], mat_c[2, 2]])
|
||||
mat = mat.transpose()
|
||||
rgb = np.dot(yuv, mat) + v
|
||||
rgb = rgb.astype(int)
|
||||
rgb = rgb.clip(0, 255)
|
||||
return rgb / 255.
|
||||
|
||||
|
||||
def read_feature_score(fp, mv_rows, mv_cols):
|
||||
line = fp.readline()
|
||||
word_ls = line.split()
|
||||
feature_score = np.array([math.log(float(v) + 1, 2) for v in word_ls])
|
||||
feature_score = feature_score.reshape(mv_rows, mv_cols)
|
||||
return feature_score
|
||||
|
||||
def read_mv_mode_arr(fp, mv_rows, mv_cols):
|
||||
line = fp.readline()
|
||||
word_ls = line.split()
|
||||
mv_mode_arr = np.array([int(v) for v in word_ls])
|
||||
mv_mode_arr = mv_mode_arr.reshape(mv_rows, mv_cols)
|
||||
return mv_mode_arr
|
||||
|
||||
|
||||
def read_frame_dpl_stats(fp):
|
||||
line = fp.readline()
|
||||
word_ls = line.split()
|
||||
frame_idx = int(word_ls[1])
|
||||
mi_rows = int(word_ls[3])
|
||||
mi_cols = int(word_ls[5])
|
||||
bs = int(word_ls[7])
|
||||
ref_frame_idx = int(word_ls[9])
|
||||
rf_idx = int(word_ls[11])
|
||||
gf_frame_offset = int(word_ls[13])
|
||||
ref_gf_frame_offset = int(word_ls[15])
|
||||
mi_size = bs / 8
|
||||
mv_ls = []
|
||||
mv_rows = int((math.ceil(mi_rows * 1. / mi_size)))
|
||||
mv_cols = int((math.ceil(mi_cols * 1. / mi_size)))
|
||||
for i in range(mv_rows * mv_cols):
|
||||
line = fp.readline()
|
||||
word_ls = line.split()
|
||||
row = int(word_ls[0]) * 8.
|
||||
col = int(word_ls[1]) * 8.
|
||||
mv_row = int(word_ls[2]) / 8.
|
||||
mv_col = int(word_ls[3]) / 8.
|
||||
mv_ls.append([col, row, mv_col, mv_row])
|
||||
mv_ls = np.array(mv_ls)
|
||||
feature_score = read_feature_score(fp, mv_rows, mv_cols)
|
||||
mv_mode_arr = read_mv_mode_arr(fp, mv_rows, mv_cols)
|
||||
img = yuv_to_rgb(read_frame(fp))
|
||||
ref = yuv_to_rgb(read_frame(fp))
|
||||
return rf_idx, frame_idx, ref_frame_idx, gf_frame_offset, ref_gf_frame_offset, mv_ls, img, ref, bs, feature_score, mv_mode_arr
|
||||
|
||||
|
||||
def read_dpl_stats_file(filename, frame_num=0):
|
||||
fp = open(filename)
|
||||
line = fp.readline()
|
||||
width = 0
|
||||
height = 0
|
||||
data_ls = []
|
||||
while (line):
|
||||
if line[0] == '=':
|
||||
data_ls.append(read_frame_dpl_stats(fp))
|
||||
line = fp.readline()
|
||||
if frame_num > 0 and len(data_ls) == frame_num:
|
||||
break
|
||||
return data_ls
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
filename = sys.argv[1]
|
||||
data_ls = read_dpl_stats_file(filename, frame_num=5)
|
||||
for rf_idx, frame_idx, ref_frame_idx, gf_frame_offset, ref_gf_frame_offset, mv_ls, img, ref, bs, feature_score, mv_mode_arr in data_ls:
|
||||
fig, axes = plt.subplots(2, 2)
|
||||
|
||||
axes[0][0].imshow(img)
|
||||
draw_mv_ls(axes[0][0], mv_ls)
|
||||
draw_pred_block_ls(axes[0][0], mv_ls, bs, mode=0)
|
||||
#axes[0].grid(color='k', linestyle='-')
|
||||
axes[0][0].set_ylim(img.shape[0], 0)
|
||||
axes[0][0].set_xlim(0, img.shape[1])
|
||||
|
||||
if ref is not None:
|
||||
axes[0][1].imshow(ref)
|
||||
draw_mv_ls(axes[0][1], mv_ls, mode=1)
|
||||
draw_pred_block_ls(axes[0][1], mv_ls, bs, mode=1)
|
||||
#axes[1].grid(color='k', linestyle='-')
|
||||
axes[0][1].set_ylim(ref.shape[0], 0)
|
||||
axes[0][1].set_xlim(0, ref.shape[1])
|
||||
|
||||
axes[1][0].imshow(feature_score)
|
||||
#feature_score_arr = feature_score.flatten()
|
||||
#feature_score_max = feature_score_arr.max()
|
||||
#feature_score_min = feature_score_arr.min()
|
||||
#step = (feature_score_max - feature_score_min) / 20.
|
||||
#feature_score_bins = np.arange(feature_score_min, feature_score_max, step)
|
||||
#axes[1][1].hist(feature_score_arr, bins=feature_score_bins)
|
||||
im = axes[1][1].imshow(mv_mode_arr)
|
||||
#axes[1][1].figure.colorbar(im, ax=axes[1][1])
|
||||
|
||||
print rf_idx, frame_idx, ref_frame_idx, gf_frame_offset, ref_gf_frame_offset, len(mv_ls)
|
||||
|
||||
flatten_mv_mode = mv_mode_arr.flatten()
|
||||
zero_mv_count = sum(flatten_mv_mode == 0);
|
||||
new_mv_count = sum(flatten_mv_mode == 1);
|
||||
ref_mv_count = sum(flatten_mv_mode == 2) + sum(flatten_mv_mode == 3);
|
||||
print zero_mv_count, new_mv_count, ref_mv_count
|
||||
plt.show()
|
||||
@@ -0,0 +1,131 @@
|
||||
## Copyright (c) 2018 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.
|
||||
##
|
||||
## Sourcing this file sets environment variables to simplify setting up
|
||||
## sanitizer builds and testing.
|
||||
|
||||
sanitizer="${1}"
|
||||
|
||||
case "${sanitizer}" in
|
||||
address) ;;
|
||||
cfi) ;;
|
||||
integer) ;;
|
||||
memory) ;;
|
||||
thread) ;;
|
||||
undefined) ;;
|
||||
clear)
|
||||
echo "Clearing environment:"
|
||||
set -x
|
||||
unset CC CXX LD AR
|
||||
unset CFLAGS CXXFLAGS LDFLAGS
|
||||
unset ASAN_OPTIONS MSAN_OPTIONS TSAN_OPTIONS UBSAN_OPTIONS
|
||||
set +x
|
||||
return
|
||||
;;
|
||||
*)
|
||||
echo "Usage: source set_analyzer_env.sh [<sanitizer>|clear]"
|
||||
echo " Supported sanitizers:"
|
||||
echo " address cfi integer memory thread undefined"
|
||||
return 1
|
||||
;;
|
||||
esac
|
||||
|
||||
if [ ! $(which clang) ]; then
|
||||
# TODO(johannkoenig): Support gcc analyzers.
|
||||
echo "ERROR: 'clang' must be in your PATH"
|
||||
return 1
|
||||
fi
|
||||
|
||||
# Warnings.
|
||||
if [ "${sanitizer}" = "undefined" -o "${sanitizer}" = "integer" ]; then
|
||||
echo "WARNING: When building the ${sanitizer} sanitizer for 32 bit targets"
|
||||
echo "you must run:"
|
||||
echo "export LDFLAGS=\"\${LDFLAGS} --rtlib=compiler-rt -lgcc_s\""
|
||||
echo "See http://llvm.org/bugs/show_bug.cgi?id=17693 for details."
|
||||
fi
|
||||
|
||||
if [ "${sanitizer}" = "undefined" ]; then
|
||||
major_version=$(clang --version | head -n 1 \
|
||||
| grep -o -E "[[:digit:]]\.[[:digit:]]\.[[:digit:]]" | cut -f1 -d.)
|
||||
if [ ${major_version} -eq 5 ]; then
|
||||
echo "WARNING: clang v5 has a problem with vp9 x86_64 high bit depth"
|
||||
echo "configurations. It can take ~40 minutes to compile"
|
||||
echo "vpx_dsp/x86/fwd_txfm_sse2.c"
|
||||
echo "clang v4 did not have this issue."
|
||||
fi
|
||||
fi
|
||||
|
||||
echo "It is recommended to configure with '--enable-debug' to improve stack"
|
||||
echo "traces. On mac builds, run 'dysmutil' on the output binaries (vpxenc,"
|
||||
echo "test_libvpx, etc) to link the stack traces to source code lines."
|
||||
|
||||
# Build configuration.
|
||||
cflags="-fsanitize=${sanitizer}"
|
||||
ldflags="-fsanitize=${sanitizer}"
|
||||
|
||||
# http://code.google.com/p/webm/issues/detail?id=570
|
||||
cflags="${cflags} -fno-strict-aliasing"
|
||||
# Useful backtraces.
|
||||
cflags="${cflags} -fno-omit-frame-pointer"
|
||||
# Exact backtraces.
|
||||
cflags="${cflags} -fno-optimize-sibling-calls"
|
||||
|
||||
if [ "${sanitizer}" = "cfi" ]; then
|
||||
# https://clang.llvm.org/docs/ControlFlowIntegrity.html
|
||||
cflags="${cflags} -fno-sanitize-trap=cfi -flto -fvisibility=hidden"
|
||||
ldflags="${ldflags} -fno-sanitize-trap=cfi -flto -fuse-ld=gold"
|
||||
export AR="llvm-ar"
|
||||
fi
|
||||
|
||||
set -x
|
||||
export CC="clang"
|
||||
export CXX="clang++"
|
||||
export LD="clang++"
|
||||
|
||||
export CFLAGS="${cflags}"
|
||||
export CXXFLAGS="${cflags}"
|
||||
export LDFLAGS="${ldflags}"
|
||||
set +x
|
||||
|
||||
# Execution configuration.
|
||||
sanitizer_options=""
|
||||
sanitizer_options="${sanitizer_options}:handle_segv=1"
|
||||
sanitizer_options="${sanitizer_options}:handle_abort=1"
|
||||
sanitizer_options="${sanitizer_options}:handle_sigfpe=1"
|
||||
sanitizer_options="${sanitizer_options}:fast_unwind_on_fatal=1"
|
||||
sanitizer_options="${sanitizer_options}:allocator_may_return_null=1"
|
||||
|
||||
case "${sanitizer}" in
|
||||
address)
|
||||
sanitizer_options="${sanitizer_options}:detect_stack_use_after_return=1"
|
||||
sanitizer_options="${sanitizer_options}:max_uar_stack_size_log=17"
|
||||
set -x
|
||||
export ASAN_OPTIONS="${sanitizer_options}"
|
||||
set +x
|
||||
;;
|
||||
cfi)
|
||||
# No environment settings
|
||||
;;
|
||||
memory)
|
||||
set -x
|
||||
export MSAN_OPTIONS="${sanitizer_options}"
|
||||
set +x
|
||||
;;
|
||||
thread)
|
||||
# The thread sanitizer uses an entirely independent set of options.
|
||||
set -x
|
||||
export TSAN_OPTIONS="halt_on_error=1"
|
||||
set +x
|
||||
;;
|
||||
undefined|integer)
|
||||
sanitizer_options="${sanitizer_options}:print_stacktrace=1"
|
||||
set -x
|
||||
export UBSAN_OPTIONS="${sanitizer_options}"
|
||||
set +x
|
||||
;;
|
||||
esac
|
||||
@@ -0,0 +1,552 @@
|
||||
/*
|
||||
* Copyright (c) 2016 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.
|
||||
*/
|
||||
|
||||
#include <assert.h>
|
||||
#include <errno.h>
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include "vpx/vpx_codec.h"
|
||||
#include "vpx/vpx_integer.h"
|
||||
#include "./y4minput.h"
|
||||
#include "vpx_dsp/ssim.h"
|
||||
#include "vpx_ports/mem.h"
|
||||
|
||||
static const int64_t cc1 = 26634; // (64^2*(.01*255)^2
|
||||
static const int64_t cc2 = 239708; // (64^2*(.03*255)^2
|
||||
static const int64_t cc1_10 = 428658; // (64^2*(.01*1023)^2
|
||||
static const int64_t cc2_10 = 3857925; // (64^2*(.03*1023)^2
|
||||
static const int64_t cc1_12 = 6868593; // (64^2*(.01*4095)^2
|
||||
static const int64_t cc2_12 = 61817334; // (64^2*(.03*4095)^2
|
||||
|
||||
#if CONFIG_VP9_HIGHBITDEPTH
|
||||
static uint64_t calc_plane_error16(uint16_t *orig, int orig_stride,
|
||||
uint16_t *recon, int recon_stride,
|
||||
unsigned int cols, unsigned int rows) {
|
||||
unsigned int row, col;
|
||||
uint64_t total_sse = 0;
|
||||
int diff;
|
||||
if (orig == NULL || recon == NULL) {
|
||||
assert(0);
|
||||
return 0;
|
||||
}
|
||||
|
||||
for (row = 0; row < rows; row++) {
|
||||
for (col = 0; col < cols; col++) {
|
||||
diff = orig[col] - recon[col];
|
||||
total_sse += diff * diff;
|
||||
}
|
||||
|
||||
orig += orig_stride;
|
||||
recon += recon_stride;
|
||||
}
|
||||
return total_sse;
|
||||
}
|
||||
#endif // CONFIG_VP9_HIGHBITDEPTH
|
||||
|
||||
static uint64_t calc_plane_error(uint8_t *orig, int orig_stride, uint8_t *recon,
|
||||
int recon_stride, unsigned int cols,
|
||||
unsigned int rows) {
|
||||
unsigned int row, col;
|
||||
uint64_t total_sse = 0;
|
||||
int diff;
|
||||
if (orig == NULL || recon == NULL) {
|
||||
assert(0);
|
||||
return 0;
|
||||
}
|
||||
|
||||
for (row = 0; row < rows; row++) {
|
||||
for (col = 0; col < cols; col++) {
|
||||
diff = orig[col] - recon[col];
|
||||
total_sse += diff * diff;
|
||||
}
|
||||
|
||||
orig += orig_stride;
|
||||
recon += recon_stride;
|
||||
}
|
||||
return total_sse;
|
||||
}
|
||||
|
||||
#define MAX_PSNR 100
|
||||
static double mse2psnr(double samples, double peak, double mse) {
|
||||
double psnr;
|
||||
|
||||
if (mse > 0.0)
|
||||
psnr = 10.0 * log10(peak * peak * samples / mse);
|
||||
else
|
||||
psnr = MAX_PSNR; // Limit to prevent / 0
|
||||
|
||||
if (psnr > MAX_PSNR) psnr = MAX_PSNR;
|
||||
|
||||
return psnr;
|
||||
}
|
||||
|
||||
typedef enum { RAW_YUV, Y4M } input_file_type;
|
||||
|
||||
typedef struct input_file {
|
||||
FILE *file;
|
||||
input_file_type type;
|
||||
unsigned char *buf;
|
||||
y4m_input y4m;
|
||||
vpx_image_t img;
|
||||
int w;
|
||||
int h;
|
||||
int bit_depth;
|
||||
int frame_size;
|
||||
} input_file_t;
|
||||
|
||||
// Open a file and determine if its y4m or raw. If y4m get the header.
|
||||
static int open_input_file(const char *file_name, input_file_t *input, int w,
|
||||
int h, int bit_depth) {
|
||||
char y4m_buf[4];
|
||||
input->w = w;
|
||||
input->h = h;
|
||||
input->bit_depth = bit_depth;
|
||||
input->type = RAW_YUV;
|
||||
input->buf = NULL;
|
||||
input->file = strcmp(file_name, "-") ? fopen(file_name, "rb") : stdin;
|
||||
if (input->file == NULL) return -1;
|
||||
if (fread(y4m_buf, 1, 4, input->file) != 4) return -1;
|
||||
if (memcmp(y4m_buf, "YUV4", 4) == 0) input->type = Y4M;
|
||||
switch (input->type) {
|
||||
case Y4M:
|
||||
y4m_input_open(&input->y4m, input->file, y4m_buf, 4, 0);
|
||||
input->w = input->y4m.pic_w;
|
||||
input->h = input->y4m.pic_h;
|
||||
input->bit_depth = input->y4m.bit_depth;
|
||||
// Y4M alloc's its own buf. Init this to avoid problems if we never
|
||||
// read frames.
|
||||
memset(&input->img, 0, sizeof(input->img));
|
||||
break;
|
||||
case RAW_YUV:
|
||||
fseek(input->file, 0, SEEK_SET);
|
||||
input->w = w;
|
||||
input->h = h;
|
||||
// handle odd frame sizes
|
||||
input->frame_size = w * h + ((w + 1) / 2) * ((h + 1) / 2) * 2;
|
||||
if (bit_depth > 8) {
|
||||
input->frame_size *= 2;
|
||||
}
|
||||
input->buf = malloc(input->frame_size);
|
||||
break;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
static void close_input_file(input_file_t *in) {
|
||||
if (in->file) fclose(in->file);
|
||||
if (in->type == Y4M) {
|
||||
vpx_img_free(&in->img);
|
||||
} else {
|
||||
free(in->buf);
|
||||
}
|
||||
}
|
||||
|
||||
static size_t read_input_file(input_file_t *in, unsigned char **y,
|
||||
unsigned char **u, unsigned char **v, int bd) {
|
||||
size_t r1 = 0;
|
||||
switch (in->type) {
|
||||
case Y4M:
|
||||
r1 = y4m_input_fetch_frame(&in->y4m, in->file, &in->img);
|
||||
*y = in->img.planes[0];
|
||||
*u = in->img.planes[1];
|
||||
*v = in->img.planes[2];
|
||||
break;
|
||||
case RAW_YUV:
|
||||
if (bd < 9) {
|
||||
r1 = fread(in->buf, in->frame_size, 1, in->file);
|
||||
*y = in->buf;
|
||||
*u = in->buf + in->w * in->h;
|
||||
*v = *u + ((1 + in->w) / 2) * ((1 + in->h) / 2);
|
||||
} else {
|
||||
r1 = fread(in->buf, in->frame_size, 1, in->file);
|
||||
*y = in->buf;
|
||||
*u = in->buf + (in->w * in->h) * 2;
|
||||
*v = *u + 2 * ((1 + in->w) / 2) * ((1 + in->h) / 2);
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
return r1;
|
||||
}
|
||||
|
||||
static void ssim_parms_8x8(const uint8_t *s, int sp, const uint8_t *r, int rp,
|
||||
uint32_t *sum_s, uint32_t *sum_r, uint32_t *sum_sq_s,
|
||||
uint32_t *sum_sq_r, uint32_t *sum_sxr) {
|
||||
int i, j;
|
||||
if (s == NULL || r == NULL || sum_s == NULL || sum_r == NULL ||
|
||||
sum_sq_s == NULL || sum_sq_r == NULL || sum_sxr == NULL) {
|
||||
assert(0);
|
||||
return;
|
||||
}
|
||||
for (i = 0; i < 8; i++, s += sp, r += rp) {
|
||||
for (j = 0; j < 8; j++) {
|
||||
*sum_s += s[j];
|
||||
*sum_r += r[j];
|
||||
*sum_sq_s += s[j] * s[j];
|
||||
*sum_sq_r += r[j] * r[j];
|
||||
*sum_sxr += s[j] * r[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if CONFIG_VP9_HIGHBITDEPTH
|
||||
static void highbd_ssim_parms_8x8(const uint16_t *s, int sp, const uint16_t *r,
|
||||
int rp, uint32_t *sum_s, uint32_t *sum_r,
|
||||
uint32_t *sum_sq_s, uint32_t *sum_sq_r,
|
||||
uint32_t *sum_sxr) {
|
||||
int i, j;
|
||||
if (s == NULL || r == NULL || sum_s == NULL || sum_r == NULL ||
|
||||
sum_sq_s == NULL || sum_sq_r == NULL || sum_sxr == NULL) {
|
||||
assert(0);
|
||||
return;
|
||||
}
|
||||
for (i = 0; i < 8; i++, s += sp, r += rp) {
|
||||
for (j = 0; j < 8; j++) {
|
||||
*sum_s += s[j];
|
||||
*sum_r += r[j];
|
||||
*sum_sq_s += s[j] * s[j];
|
||||
*sum_sq_r += r[j] * r[j];
|
||||
*sum_sxr += s[j] * r[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif // CONFIG_VP9_HIGHBITDEPTH
|
||||
|
||||
static double similarity(uint32_t sum_s, uint32_t sum_r, uint32_t sum_sq_s,
|
||||
uint32_t sum_sq_r, uint32_t sum_sxr, int count,
|
||||
uint32_t bd) {
|
||||
double ssim_n, ssim_d;
|
||||
int64_t c1 = 0, c2 = 0;
|
||||
if (bd == 8) {
|
||||
// scale the constants by number of pixels
|
||||
c1 = (cc1 * count * count) >> 12;
|
||||
c2 = (cc2 * count * count) >> 12;
|
||||
} else if (bd == 10) {
|
||||
c1 = (cc1_10 * count * count) >> 12;
|
||||
c2 = (cc2_10 * count * count) >> 12;
|
||||
} else if (bd == 12) {
|
||||
c1 = (cc1_12 * count * count) >> 12;
|
||||
c2 = (cc2_12 * count * count) >> 12;
|
||||
} else {
|
||||
assert(0);
|
||||
}
|
||||
|
||||
ssim_n = (2.0 * sum_s * sum_r + c1) *
|
||||
(2.0 * count * sum_sxr - 2.0 * sum_s * sum_r + c2);
|
||||
|
||||
ssim_d = ((double)sum_s * sum_s + (double)sum_r * sum_r + c1) *
|
||||
((double)count * sum_sq_s - (double)sum_s * sum_s +
|
||||
(double)count * sum_sq_r - (double)sum_r * sum_r + c2);
|
||||
|
||||
return ssim_n / ssim_d;
|
||||
}
|
||||
|
||||
static double ssim_8x8(const uint8_t *s, int sp, const uint8_t *r, int rp) {
|
||||
uint32_t sum_s = 0, sum_r = 0, sum_sq_s = 0, sum_sq_r = 0, sum_sxr = 0;
|
||||
ssim_parms_8x8(s, sp, r, rp, &sum_s, &sum_r, &sum_sq_s, &sum_sq_r, &sum_sxr);
|
||||
return similarity(sum_s, sum_r, sum_sq_s, sum_sq_r, sum_sxr, 64, 8);
|
||||
}
|
||||
|
||||
#if CONFIG_VP9_HIGHBITDEPTH
|
||||
static double highbd_ssim_8x8(const uint16_t *s, int sp, const uint16_t *r,
|
||||
int rp, uint32_t bd) {
|
||||
uint32_t sum_s = 0, sum_r = 0, sum_sq_s = 0, sum_sq_r = 0, sum_sxr = 0;
|
||||
highbd_ssim_parms_8x8(s, sp, r, rp, &sum_s, &sum_r, &sum_sq_s, &sum_sq_r,
|
||||
&sum_sxr);
|
||||
return similarity(sum_s, sum_r, sum_sq_s, sum_sq_r, sum_sxr, 64, bd);
|
||||
}
|
||||
#endif // CONFIG_VP9_HIGHBITDEPTH
|
||||
|
||||
// We are using a 8x8 moving window with starting location of each 8x8 window
|
||||
// on the 4x4 pixel grid. Such arrangement allows the windows to overlap
|
||||
// block boundaries to penalize blocking artifacts.
|
||||
static double ssim2(const uint8_t *img1, const uint8_t *img2, int stride_img1,
|
||||
int stride_img2, int width, int height) {
|
||||
int i, j;
|
||||
int samples = 0;
|
||||
double ssim_total = 0;
|
||||
|
||||
// sample point start with each 4x4 location
|
||||
for (i = 0; i <= height - 8;
|
||||
i += 4, img1 += stride_img1 * 4, img2 += stride_img2 * 4) {
|
||||
for (j = 0; j <= width - 8; j += 4) {
|
||||
double v = ssim_8x8(img1 + j, stride_img1, img2 + j, stride_img2);
|
||||
ssim_total += v;
|
||||
samples++;
|
||||
}
|
||||
}
|
||||
ssim_total /= samples;
|
||||
return ssim_total;
|
||||
}
|
||||
|
||||
#if CONFIG_VP9_HIGHBITDEPTH
|
||||
static double highbd_ssim2(const uint8_t *img1, const uint8_t *img2,
|
||||
int stride_img1, int stride_img2, int width,
|
||||
int height, uint32_t bd) {
|
||||
int i, j;
|
||||
int samples = 0;
|
||||
double ssim_total = 0;
|
||||
|
||||
// sample point start with each 4x4 location
|
||||
for (i = 0; i <= height - 8;
|
||||
i += 4, img1 += stride_img1 * 4, img2 += stride_img2 * 4) {
|
||||
for (j = 0; j <= width - 8; j += 4) {
|
||||
double v =
|
||||
highbd_ssim_8x8(CONVERT_TO_SHORTPTR(img1 + j), stride_img1,
|
||||
CONVERT_TO_SHORTPTR(img2 + j), stride_img2, bd);
|
||||
ssim_total += v;
|
||||
samples++;
|
||||
}
|
||||
}
|
||||
ssim_total /= samples;
|
||||
return ssim_total;
|
||||
}
|
||||
#endif // CONFIG_VP9_HIGHBITDEPTH
|
||||
|
||||
int main(int argc, char *argv[]) {
|
||||
FILE *framestats = NULL;
|
||||
int bit_depth = 8;
|
||||
int w = 0, h = 0, tl_skip = 0, tl_skips_remaining = 0;
|
||||
double ssimavg = 0, ssimyavg = 0, ssimuavg = 0, ssimvavg = 0;
|
||||
double psnrglb = 0, psnryglb = 0, psnruglb = 0, psnrvglb = 0;
|
||||
double psnravg = 0, psnryavg = 0, psnruavg = 0, psnrvavg = 0;
|
||||
double *ssimy = NULL, *ssimu = NULL, *ssimv = NULL;
|
||||
uint64_t *psnry = NULL, *psnru = NULL, *psnrv = NULL;
|
||||
size_t i, n_frames = 0, allocated_frames = 0;
|
||||
int return_value = 0;
|
||||
input_file_t in[2];
|
||||
double peak = 255.0;
|
||||
|
||||
memset(in, 0, sizeof(in));
|
||||
|
||||
if (argc < 2) {
|
||||
fprintf(stderr,
|
||||
"Usage: %s file1.{yuv|y4m} file2.{yuv|y4m}"
|
||||
"[WxH tl_skip={0,1,3} frame_stats_file bits]\n",
|
||||
argv[0]);
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (argc > 3) {
|
||||
sscanf(argv[3], "%dx%d", &w, &h);
|
||||
}
|
||||
|
||||
if (argc > 6) {
|
||||
sscanf(argv[6], "%d", &bit_depth);
|
||||
}
|
||||
|
||||
if (open_input_file(argv[1], &in[0], w, h, bit_depth) < 0) {
|
||||
fprintf(stderr, "File %s can't be opened or parsed!\n", argv[1]);
|
||||
goto clean_up;
|
||||
}
|
||||
|
||||
if (w == 0 && h == 0) {
|
||||
// If a y4m is the first file and w, h is not set grab from first file.
|
||||
w = in[0].w;
|
||||
h = in[0].h;
|
||||
bit_depth = in[0].bit_depth;
|
||||
}
|
||||
if (bit_depth == 10) peak = 1023.0;
|
||||
|
||||
if (bit_depth == 12) peak = 4095.0;
|
||||
|
||||
if (open_input_file(argv[2], &in[1], w, h, bit_depth) < 0) {
|
||||
fprintf(stderr, "File %s can't be opened or parsed!\n", argv[2]);
|
||||
goto clean_up;
|
||||
}
|
||||
|
||||
if (in[0].w != in[1].w || in[0].h != in[1].h || in[0].w != w ||
|
||||
in[0].h != h || w == 0 || h == 0) {
|
||||
fprintf(stderr,
|
||||
"Failing: Image dimensions don't match or are unspecified!\n");
|
||||
return_value = 1;
|
||||
goto clean_up;
|
||||
}
|
||||
|
||||
if (in[0].bit_depth != in[1].bit_depth) {
|
||||
fprintf(stderr,
|
||||
"Failing: Image bit depths don't match or are unspecified!\n");
|
||||
return_value = 1;
|
||||
goto clean_up;
|
||||
}
|
||||
|
||||
bit_depth = in[0].bit_depth;
|
||||
|
||||
// Number of frames to skip from file1.yuv for every frame used. Normal
|
||||
// values 0, 1 and 3 correspond to TL2, TL1 and TL0 respectively for a 3TL
|
||||
// encoding in mode 10. 7 would be reasonable for comparing TL0 of a 4-layer
|
||||
// encoding.
|
||||
if (argc > 4) {
|
||||
sscanf(argv[4], "%d", &tl_skip);
|
||||
if (argc > 5) {
|
||||
framestats = fopen(argv[5], "w");
|
||||
if (!framestats) {
|
||||
fprintf(stderr, "Could not open \"%s\" for writing: %s\n", argv[5],
|
||||
strerror(errno));
|
||||
return_value = 1;
|
||||
goto clean_up;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
while (1) {
|
||||
size_t r1, r2;
|
||||
unsigned char *y[2], *u[2], *v[2];
|
||||
|
||||
r1 = read_input_file(&in[0], &y[0], &u[0], &v[0], bit_depth);
|
||||
|
||||
if (r1) {
|
||||
// Reading parts of file1.yuv that were not used in temporal layer.
|
||||
if (tl_skips_remaining > 0) {
|
||||
--tl_skips_remaining;
|
||||
continue;
|
||||
}
|
||||
// Use frame, but skip |tl_skip| after it.
|
||||
tl_skips_remaining = tl_skip;
|
||||
}
|
||||
|
||||
r2 = read_input_file(&in[1], &y[1], &u[1], &v[1], bit_depth);
|
||||
|
||||
if (r1 && r2 && r1 != r2) {
|
||||
fprintf(stderr, "Failed to read data: %s [%d/%d]\n", strerror(errno),
|
||||
(int)r1, (int)r2);
|
||||
return_value = 1;
|
||||
goto clean_up;
|
||||
} else if (r1 == 0 || r2 == 0) {
|
||||
break;
|
||||
}
|
||||
#if CONFIG_VP9_HIGHBITDEPTH
|
||||
#define psnr_and_ssim(ssim, psnr, buf0, buf1, w, h) \
|
||||
if (bit_depth < 9) { \
|
||||
ssim = ssim2(buf0, buf1, w, w, w, h); \
|
||||
psnr = calc_plane_error(buf0, w, buf1, w, w, h); \
|
||||
} else { \
|
||||
ssim = highbd_ssim2(CONVERT_TO_BYTEPTR(buf0), CONVERT_TO_BYTEPTR(buf1), w, \
|
||||
w, w, h, bit_depth); \
|
||||
psnr = calc_plane_error16(CAST_TO_SHORTPTR(buf0), w, \
|
||||
CAST_TO_SHORTPTR(buf1), w, w, h); \
|
||||
}
|
||||
#else
|
||||
#define psnr_and_ssim(ssim, psnr, buf0, buf1, w, h) \
|
||||
ssim = ssim2(buf0, buf1, w, w, w, h); \
|
||||
psnr = calc_plane_error(buf0, w, buf1, w, w, h);
|
||||
#endif // CONFIG_VP9_HIGHBITDEPTH
|
||||
|
||||
if (n_frames == allocated_frames) {
|
||||
allocated_frames = allocated_frames == 0 ? 1024 : allocated_frames * 2;
|
||||
ssimy = realloc(ssimy, allocated_frames * sizeof(*ssimy));
|
||||
ssimu = realloc(ssimu, allocated_frames * sizeof(*ssimu));
|
||||
ssimv = realloc(ssimv, allocated_frames * sizeof(*ssimv));
|
||||
psnry = realloc(psnry, allocated_frames * sizeof(*psnry));
|
||||
psnru = realloc(psnru, allocated_frames * sizeof(*psnru));
|
||||
psnrv = realloc(psnrv, allocated_frames * sizeof(*psnrv));
|
||||
}
|
||||
psnr_and_ssim(ssimy[n_frames], psnry[n_frames], y[0], y[1], w, h);
|
||||
psnr_and_ssim(ssimu[n_frames], psnru[n_frames], u[0], u[1], (w + 1) / 2,
|
||||
(h + 1) / 2);
|
||||
psnr_and_ssim(ssimv[n_frames], psnrv[n_frames], v[0], v[1], (w + 1) / 2,
|
||||
(h + 1) / 2);
|
||||
|
||||
n_frames++;
|
||||
}
|
||||
|
||||
if (framestats) {
|
||||
fprintf(framestats,
|
||||
"ssim,ssim-y,ssim-u,ssim-v,psnr,psnr-y,psnr-u,psnr-v\n");
|
||||
}
|
||||
|
||||
for (i = 0; i < n_frames; ++i) {
|
||||
double frame_ssim;
|
||||
double frame_psnr, frame_psnry, frame_psnru, frame_psnrv;
|
||||
|
||||
frame_ssim = 0.8 * ssimy[i] + 0.1 * (ssimu[i] + ssimv[i]);
|
||||
ssimavg += frame_ssim;
|
||||
ssimyavg += ssimy[i];
|
||||
ssimuavg += ssimu[i];
|
||||
ssimvavg += ssimv[i];
|
||||
|
||||
frame_psnr =
|
||||
mse2psnr(w * h * 6 / 4, peak, (double)psnry[i] + psnru[i] + psnrv[i]);
|
||||
frame_psnry = mse2psnr(w * h * 4 / 4, peak, (double)psnry[i]);
|
||||
frame_psnru = mse2psnr(w * h * 1 / 4, peak, (double)psnru[i]);
|
||||
frame_psnrv = mse2psnr(w * h * 1 / 4, peak, (double)psnrv[i]);
|
||||
|
||||
psnravg += frame_psnr;
|
||||
psnryavg += frame_psnry;
|
||||
psnruavg += frame_psnru;
|
||||
psnrvavg += frame_psnrv;
|
||||
|
||||
psnryglb += psnry[i];
|
||||
psnruglb += psnru[i];
|
||||
psnrvglb += psnrv[i];
|
||||
|
||||
if (framestats) {
|
||||
fprintf(framestats, "%lf,%lf,%lf,%lf,%lf,%lf,%lf,%lf\n", frame_ssim,
|
||||
ssimy[i], ssimu[i], ssimv[i], frame_psnr, frame_psnry,
|
||||
frame_psnru, frame_psnrv);
|
||||
}
|
||||
}
|
||||
|
||||
ssimavg /= n_frames;
|
||||
ssimyavg /= n_frames;
|
||||
ssimuavg /= n_frames;
|
||||
ssimvavg /= n_frames;
|
||||
|
||||
printf("VpxSSIM: %lf\n", 100 * pow(ssimavg, 8.0));
|
||||
printf("SSIM: %lf\n", ssimavg);
|
||||
printf("SSIM-Y: %lf\n", ssimyavg);
|
||||
printf("SSIM-U: %lf\n", ssimuavg);
|
||||
printf("SSIM-V: %lf\n", ssimvavg);
|
||||
puts("");
|
||||
|
||||
psnravg /= n_frames;
|
||||
psnryavg /= n_frames;
|
||||
psnruavg /= n_frames;
|
||||
psnrvavg /= n_frames;
|
||||
|
||||
printf("AvgPSNR: %lf\n", psnravg);
|
||||
printf("AvgPSNR-Y: %lf\n", psnryavg);
|
||||
printf("AvgPSNR-U: %lf\n", psnruavg);
|
||||
printf("AvgPSNR-V: %lf\n", psnrvavg);
|
||||
puts("");
|
||||
|
||||
psnrglb = psnryglb + psnruglb + psnrvglb;
|
||||
psnrglb = mse2psnr((double)n_frames * w * h * 6 / 4, peak, psnrglb);
|
||||
psnryglb = mse2psnr((double)n_frames * w * h * 4 / 4, peak, psnryglb);
|
||||
psnruglb = mse2psnr((double)n_frames * w * h * 1 / 4, peak, psnruglb);
|
||||
psnrvglb = mse2psnr((double)n_frames * w * h * 1 / 4, peak, psnrvglb);
|
||||
|
||||
printf("GlbPSNR: %lf\n", psnrglb);
|
||||
printf("GlbPSNR-Y: %lf\n", psnryglb);
|
||||
printf("GlbPSNR-U: %lf\n", psnruglb);
|
||||
printf("GlbPSNR-V: %lf\n", psnrvglb);
|
||||
puts("");
|
||||
|
||||
printf("Nframes: %d\n", (int)n_frames);
|
||||
|
||||
clean_up:
|
||||
|
||||
close_input_file(&in[0]);
|
||||
close_input_file(&in[1]);
|
||||
|
||||
if (framestats) fclose(framestats);
|
||||
|
||||
free(ssimy);
|
||||
free(ssimu);
|
||||
free(ssimv);
|
||||
|
||||
free(psnry);
|
||||
free(psnru);
|
||||
free(psnrv);
|
||||
|
||||
return return_value;
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
#!/usr/bin/env python
|
||||
## Copyright (c) 2012 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.
|
||||
##
|
||||
"""Wraps paragraphs of text, preserving manual formatting
|
||||
|
||||
This is like fold(1), but has the special convention of not modifying lines
|
||||
that start with whitespace. This allows you to intersperse blocks with
|
||||
special formatting, like code blocks, with written prose. The prose will
|
||||
be wordwrapped, and the manual formatting will be preserved.
|
||||
|
||||
* This won't handle the case of a bulleted (or ordered) list specially, so
|
||||
manual wrapping must be done.
|
||||
|
||||
Occasionally it's useful to put something with explicit formatting that
|
||||
doesn't look at all like a block of text inline.
|
||||
|
||||
indicator = has_leading_whitespace(line);
|
||||
if (indicator)
|
||||
preserve_formatting(line);
|
||||
|
||||
The intent is that this docstring would make it through the transform
|
||||
and still be legible and presented as it is in the source. If additional
|
||||
cases are handled, update this doc to describe the effect.
|
||||
"""
|
||||
|
||||
__author__ = "jkoleszar@google.com"
|
||||
import textwrap
|
||||
import sys
|
||||
|
||||
def wrap(text):
|
||||
if text:
|
||||
return textwrap.fill(text, break_long_words=False) + '\n'
|
||||
return ""
|
||||
|
||||
|
||||
def main(fileobj):
|
||||
text = ""
|
||||
output = ""
|
||||
while True:
|
||||
line = fileobj.readline()
|
||||
if not line:
|
||||
break
|
||||
|
||||
if line.lstrip() == line:
|
||||
text += line
|
||||
else:
|
||||
output += wrap(text)
|
||||
text=""
|
||||
output += line
|
||||
output += wrap(text)
|
||||
|
||||
# Replace the file or write to stdout.
|
||||
if fileobj == sys.stdin:
|
||||
fileobj = sys.stdout
|
||||
else:
|
||||
fileobj.seek(0)
|
||||
fileobj.truncate(0)
|
||||
fileobj.write(output)
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) > 1:
|
||||
main(open(sys.argv[1], "r+"))
|
||||
else:
|
||||
main(sys.stdin)
|
||||
Reference in New Issue
Block a user