1## Copyright (c) 2020 The WebM project authors. All Rights Reserved. 2## 3## Use of this source code is governed by a BSD-style license 4## that can be found in the LICENSE file in the root of the source 5## tree. An additional intellectual property rights grant can be found 6## in the file PATENTS. All contributing project authors may 7## be found in the AUTHORS file in the root of the source tree. 8## 9 10import sys 11import matplotlib.pyplot as plt 12from matplotlib.collections import LineCollection 13from matplotlib import colors as mcolors 14import numpy as np 15import math 16 17 18def draw_mv_ls(axis, mv_ls, mode=0): 19 colors = np.array([(1., 0., 0., 1.)]) 20 segs = np.array([ 21 np.array([[ptr[0], ptr[1]], [ptr[0] + ptr[2], ptr[1] + ptr[3]]]) 22 for ptr in mv_ls 23 ]) 24 line_segments = LineCollection( 25 segs, linewidths=(1.,), colors=colors, linestyle='solid') 26 axis.add_collection(line_segments) 27 if mode == 0: 28 axis.scatter(mv_ls[:, 0], mv_ls[:, 1], s=2, c='b') 29 else: 30 axis.scatter( 31 mv_ls[:, 0] + mv_ls[:, 2], mv_ls[:, 1] + mv_ls[:, 3], s=2, c='b') 32 33 34def draw_pred_block_ls(axis, mv_ls, bs, mode=0): 35 colors = np.array([(0., 0., 0., 1.)]) 36 segs = [] 37 for ptr in mv_ls: 38 if mode == 0: 39 x = ptr[0] 40 y = ptr[1] 41 else: 42 x = ptr[0] + ptr[2] 43 y = ptr[1] + ptr[3] 44 x_ls = [x, x + bs, x + bs, x, x] 45 y_ls = [y, y, y + bs, y + bs, y] 46 47 segs.append(np.column_stack([x_ls, y_ls])) 48 line_segments = LineCollection( 49 segs, linewidths=(.5,), colors=colors, linestyle='solid') 50 axis.add_collection(line_segments) 51 52 53def read_frame(fp, no_swap=0): 54 plane = [None, None, None] 55 for i in range(3): 56 line = fp.readline() 57 word_ls = line.split() 58 word_ls = [int(item) for item in word_ls] 59 rows = word_ls[0] 60 cols = word_ls[1] 61 62 line = fp.readline() 63 word_ls = line.split() 64 word_ls = [int(item) for item in word_ls] 65 66 plane[i] = np.array(word_ls).reshape(rows, cols) 67 if i > 0: 68 plane[i] = plane[i].repeat(2, axis=0).repeat(2, axis=1) 69 plane = np.array(plane) 70 if no_swap == 0: 71 plane = np.swapaxes(np.swapaxes(plane, 0, 1), 1, 2) 72 return plane 73 74 75def yuv_to_rgb(yuv): 76 #mat = np.array([ 77 # [1.164, 0 , 1.596 ], 78 # [1.164, -0.391, -0.813], 79 # [1.164, 2.018 , 0 ] ] 80 # ) 81 #c = np.array([[ -16 , -16 , -16 ], 82 # [ 0 , -128, -128 ], 83 # [ -128, -128, 0 ]]) 84 85 mat = np.array([[1, 0, 1.4075], [1, -0.3445, -0.7169], [1, 1.7790, 0]]) 86 c = np.array([[0, 0, 0], [0, -128, -128], [-128, -128, 0]]) 87 mat_c = np.dot(mat, c) 88 v = np.array([mat_c[0, 0], mat_c[1, 1], mat_c[2, 2]]) 89 mat = mat.transpose() 90 rgb = np.dot(yuv, mat) + v 91 rgb = rgb.astype(int) 92 rgb = rgb.clip(0, 255) 93 return rgb / 255. 94 95 96def read_feature_score(fp, mv_rows, mv_cols): 97 line = fp.readline() 98 word_ls = line.split() 99 feature_score = np.array([math.log(float(v) + 1, 2) for v in word_ls]) 100 feature_score = feature_score.reshape(mv_rows, mv_cols) 101 return feature_score 102 103def read_mv_mode_arr(fp, mv_rows, mv_cols): 104 line = fp.readline() 105 word_ls = line.split() 106 mv_mode_arr = np.array([int(v) for v in word_ls]) 107 mv_mode_arr = mv_mode_arr.reshape(mv_rows, mv_cols) 108 return mv_mode_arr 109 110 111def read_frame_dpl_stats(fp): 112 line = fp.readline() 113 word_ls = line.split() 114 frame_idx = int(word_ls[1]) 115 mi_rows = int(word_ls[3]) 116 mi_cols = int(word_ls[5]) 117 bs = int(word_ls[7]) 118 ref_frame_idx = int(word_ls[9]) 119 rf_idx = int(word_ls[11]) 120 gf_frame_offset = int(word_ls[13]) 121 ref_gf_frame_offset = int(word_ls[15]) 122 mi_size = bs / 8 123 mv_ls = [] 124 mv_rows = int((math.ceil(mi_rows * 1. / mi_size))) 125 mv_cols = int((math.ceil(mi_cols * 1. / mi_size))) 126 for i in range(mv_rows * mv_cols): 127 line = fp.readline() 128 word_ls = line.split() 129 row = int(word_ls[0]) * 8. 130 col = int(word_ls[1]) * 8. 131 mv_row = int(word_ls[2]) / 8. 132 mv_col = int(word_ls[3]) / 8. 133 mv_ls.append([col, row, mv_col, mv_row]) 134 mv_ls = np.array(mv_ls) 135 feature_score = read_feature_score(fp, mv_rows, mv_cols) 136 mv_mode_arr = read_mv_mode_arr(fp, mv_rows, mv_cols) 137 img = yuv_to_rgb(read_frame(fp)) 138 ref = yuv_to_rgb(read_frame(fp)) 139 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 140 141 142def read_dpl_stats_file(filename, frame_num=0): 143 fp = open(filename) 144 line = fp.readline() 145 width = 0 146 height = 0 147 data_ls = [] 148 while (line): 149 if line[0] == '=': 150 data_ls.append(read_frame_dpl_stats(fp)) 151 line = fp.readline() 152 if frame_num > 0 and len(data_ls) == frame_num: 153 break 154 return data_ls 155 156 157if __name__ == '__main__': 158 filename = sys.argv[1] 159 data_ls = read_dpl_stats_file(filename, frame_num=5) 160 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: 161 fig, axes = plt.subplots(2, 2) 162 163 axes[0][0].imshow(img) 164 draw_mv_ls(axes[0][0], mv_ls) 165 draw_pred_block_ls(axes[0][0], mv_ls, bs, mode=0) 166 #axes[0].grid(color='k', linestyle='-') 167 axes[0][0].set_ylim(img.shape[0], 0) 168 axes[0][0].set_xlim(0, img.shape[1]) 169 170 if ref is not None: 171 axes[0][1].imshow(ref) 172 draw_mv_ls(axes[0][1], mv_ls, mode=1) 173 draw_pred_block_ls(axes[0][1], mv_ls, bs, mode=1) 174 #axes[1].grid(color='k', linestyle='-') 175 axes[0][1].set_ylim(ref.shape[0], 0) 176 axes[0][1].set_xlim(0, ref.shape[1]) 177 178 axes[1][0].imshow(feature_score) 179 #feature_score_arr = feature_score.flatten() 180 #feature_score_max = feature_score_arr.max() 181 #feature_score_min = feature_score_arr.min() 182 #step = (feature_score_max - feature_score_min) / 20. 183 #feature_score_bins = np.arange(feature_score_min, feature_score_max, step) 184 #axes[1][1].hist(feature_score_arr, bins=feature_score_bins) 185 im = axes[1][1].imshow(mv_mode_arr) 186 #axes[1][1].figure.colorbar(im, ax=axes[1][1]) 187 188 print rf_idx, frame_idx, ref_frame_idx, gf_frame_offset, ref_gf_frame_offset, len(mv_ls) 189 190 flatten_mv_mode = mv_mode_arr.flatten() 191 zero_mv_count = sum(flatten_mv_mode == 0); 192 new_mv_count = sum(flatten_mv_mode == 1); 193 ref_mv_count = sum(flatten_mv_mode == 2) + sum(flatten_mv_mode == 3); 194 print zero_mv_count, new_mv_count, ref_mv_count 195 plt.show() 196