Include gf_frame_offset in dump_tpl_stats

Change-Id: If868f67ccc1c73189bc4c139a807d7341e59b668

1 files changed, 6 insertions, 3 deletions

diff --git a/tools/non_greedy_mv/non_greedy_mv.py b/tools/non_greedy_mv/non_greedy_mv.py
index 3409e70fc..431f2434b 100644
--- a/tools/non_greedy_mv/non_greedy_mv.py
+++ b/tools/non_greedy_mv/non_greedy_mv.py
@@ -100,6 +100,9 @@ def read_frame_dpl_stats(fp):
   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)))
@@ -116,7 +119,7 @@ def read_frame_dpl_stats(fp):
   feature_score = read_feature_score(fp, mv_rows, mv_cols)
   img = yuv_to_rgb(read_frame(fp))
   ref = yuv_to_rgb(read_frame(fp))
-  return frame_idx, ref_frame_idx, mv_ls, img, ref, bs, feature_score
+  return rf_idx, frame_idx, ref_frame_idx, gf_frame_offset, ref_gf_frame_offset, mv_ls, img, ref, bs, feature_score
 
 
 def read_dpl_stats_file(filename, frame_num=0):
@@ -137,7 +140,7 @@ def read_dpl_stats_file(filename, frame_num=0):
 if __name__ == '__main__':
   filename = sys.argv[1]
   data_ls = read_dpl_stats_file(filename, frame_num=5)
-  for frame_idx, ref_frame_idx, mv_ls, img, ref, bs, feature_score in data_ls:
+  for rf_idx, frame_idx, ref_frame_idx, gf_frame_offset, ref_gf_frame_offset, mv_ls, img, ref, bs, feature_score in data_ls:
     fig, axes = plt.subplots(2, 2)
 
     axes[0][0].imshow(img)
@@ -163,5 +166,5 @@ if __name__ == '__main__':
     feature_score_bins = np.arange(feature_score_min, feature_score_max, step)
     axes[1][1].hist(feature_score_arr, bins=feature_score_bins)
 
-    print frame_idx, ref_frame_idx, len(mv_ls)
+    print rf_idx, frame_idx, ref_frame_idx, gf_frame_offset, ref_gf_frame_offset, len(mv_ls)
     plt.show()