tools/non_greedy_mv/non_greedy_mv.py

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