diff options
Diffstat (limited to 'tools/non_greedy_mv/non_greedy_mv.py')
-rw-r--r-- | tools/non_greedy_mv/non_greedy_mv.py | 195 |
1 files changed, 195 insertions, 0 deletions
diff --git a/tools/non_greedy_mv/non_greedy_mv.py b/tools/non_greedy_mv/non_greedy_mv.py new file mode 100644 index 000000000..a46b7e760 --- /dev/null +++ b/tools/non_greedy_mv/non_greedy_mv.py @@ -0,0 +1,195 @@ +## Copyright (c) 2020 The WebM project authors. All Rights Reserved. +## +## Use of this source code is governed by a BSD-style license +## that can be found in the LICENSE file in the root of the source +## tree. An additional intellectual property rights grant can be found +## in the file PATENTS. All contributing project authors may +## be found in the AUTHORS file in the root of the source tree. +## + +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() |