Include gf_frame_offset in dump_tpl_stats

Change-Id: If868f67ccc1c73189bc4c139a807d7341e59b668

2 files changed, 16 insertions, 7 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()
diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index 65113a660..f70ded054 100644
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -6746,6 +6746,7 @@ static void dump_frame_buf(const YV12_BUFFER_CONFIG *frame_buf) {
 }
 
 static void dump_tpl_stats(const VP9_COMP *cpi, int tpl_group_frames,
+                           const GF_GROUP *gf_group,
                            const GF_PICTURE *gf_picture, BLOCK_SIZE bsize) {
   int frame_idx;
   const VP9_COMMON *cm = &cpi->common;
@@ -6760,10 +6761,15 @@ static void dump_tpl_stats(const VP9_COMP *cpi, int tpl_group_frames,
       ref_frame_idx = gf_picture[frame_idx].ref_frame[rf_idx];
       if (ref_frame_idx != -1) {
         YV12_BUFFER_CONFIG *ref_frame_buf = gf_picture[ref_frame_idx].frame;
+        const int gf_frame_offset = gf_group->frame_gop_index[frame_idx];
+        const int ref_gf_frame_offset =
+            gf_group->frame_gop_index[ref_frame_idx];
         printf("=\n");
-        printf("frame_idx %d mi_rows %d mi_cols %d bsize %d ref_frame_idx %d\n",
-               frame_idx, cm->mi_rows, cm->mi_cols, mi_width * MI_SIZE,
-               ref_frame_idx);
+        printf(
+            "frame_idx %d mi_rows %d mi_cols %d bsize %d ref_frame_idx %d "
+            "rf_idx %d gf_frame_offset %d ref_gf_frame_offset %d\n",
+            frame_idx, cm->mi_rows, cm->mi_cols, mi_width * MI_SIZE,
+            ref_frame_idx, rf_idx, gf_frame_offset, ref_gf_frame_offset);
         for (mi_row = 0; mi_row < cm->mi_rows; ++mi_row) {
           for (mi_col = 0; mi_col < cm->mi_cols; ++mi_col) {
             if ((mi_row % mi_height) == 0 && (mi_col % mi_width) == 0) {
@@ -6890,7 +6896,7 @@ static void setup_tpl_stats(VP9_COMP *cpi) {
 #if CONFIG_NON_GREEDY_MV
   cpi->tpl_ready = 1;
 #if DUMP_TPL_STATS
-  dump_tpl_stats(cpi, tpl_group_frames, gf_picture, cpi->tpl_bsize);
+  dump_tpl_stats(cpi, tpl_group_frames, gf_group, gf_picture, cpi->tpl_bsize);
 #endif  // DUMP_TPL_STATS
 #endif  // CONFIG_NON_GREEDY_MV
 }