Add extra check / unit test to SetExternalGroupOfPicturesMap()

Let SetExternalGroupOfPicturesMap() modify the gop_map_ to satisfy
the following constraints.
1) Each key frame position should be at the start of a gop.
2) The last gop should not use an alt ref.

Add unit test for SetExternalGroupOfPicturesMap()

Change-Id: Iee9bd238ad0fc5c2ccbf2fbd065a280c854cd718

3 files changed, 111 insertions, 5 deletions

diff --git a/test/simple_encode_test.cc b/test/simple_encode_test.cc
index cbcae46be..eddd1c59b 100644
--- a/test/simple_encode_test.cc
+++ b/test/simple_encode_test.cc
@@ -319,7 +319,7 @@ TEST(SimpleEncode, EncodeConsistencyTest3) {
 // Get QPs and arf locations from the first encode.
 // Set external arfs and QPs for the second encode.
 // Expect to get matched results.
-TEST(SimpleEncode, EncodeConsistencySetExternalGroupOfPicture) {
+TEST(SimpleEncode, EncodeConsistencySetExternalGroupOfPicturesMap) {
   std::vector<int> quantize_index_list;
   std::vector<uint64_t> ref_sse_list;
   std::vector<double> ref_psnr_list;
@@ -361,7 +361,7 @@ TEST(SimpleEncode, EncodeConsistencySetExternalGroupOfPicture) {
     SimpleEncode simple_encode(w, h, frame_rate_num, frame_rate_den,
                                target_bitrate, num_frames, infile_path);
     simple_encode.ComputeFirstPassStats();
-    simple_encode.SetExternalGroupOfPicture(gop_map);
+    simple_encode.SetExternalGroupOfPicturesMap(gop_map);
     const int num_coding_frames = simple_encode.GetCodingFrameNum();
     EXPECT_EQ(static_cast<size_t>(num_coding_frames),
               quantize_index_list.size());
@@ -379,6 +379,73 @@ TEST(SimpleEncode, EncodeConsistencySetExternalGroupOfPicture) {
   }
 }
 
+TEST(SimpleEncode, SetExternalGroupOfPicturesMap) {
+  SimpleEncode simple_encode(w, h, frame_rate_num, frame_rate_den,
+                             target_bitrate, num_frames, infile_path);
+  simple_encode.ComputeFirstPassStats();
+
+  std::vector<int> gop_map(num_frames, 0);
+
+  // Should be the first gop group.
+  gop_map[0] = 0;
+
+  // Second gop group with an alt ref.
+  gop_map[5] |= kGopMapFlagStart | kGopMapFlagUseAltRef;
+
+  // Third gop group without an alt ref.
+  gop_map[10] |= kGopMapFlagStart;
+
+  // Last gop group.
+  gop_map[14] |= kGopMapFlagStart | kGopMapFlagUseAltRef;
+
+  simple_encode.SetExternalGroupOfPicturesMap(gop_map);
+
+  std::vector<int> observed_gop_map =
+      simple_encode.ObserveExternalGroupOfPicturesMap();
+
+  // First gop group.
+  // There is always a key frame at show_idx 0 and key frame should always be
+  // the start of a gop. We expect ObserveExternalGroupOfPicturesMap() will
+  // insert an extra gop start here.
+  EXPECT_EQ(observed_gop_map[0], kGopMapFlagStart | kGopMapFlagUseAltRef);
+
+  // Second gop group with an alt ref.
+  EXPECT_EQ(observed_gop_map[5], kGopMapFlagStart | kGopMapFlagUseAltRef);
+
+  // Third gop group without an alt ref.
+  EXPECT_EQ(observed_gop_map[10], kGopMapFlagStart);
+
+  // Last gop group. The last gop is not supposed to use an alt ref. We expect
+  // ObserveExternalGroupOfPicturesMap() will remove the alt ref flag here.
+  EXPECT_EQ(observed_gop_map[14], kGopMapFlagStart);
+
+  int ref_gop_show_frame_count_list[4] = { 5, 5, 4, 3 };
+  size_t ref_gop_coded_frame_count_list[4] = { 6, 6, 4, 3 };
+  int gop_count = 0;
+
+  simple_encode.StartEncode();
+  int coded_show_frame_count = 0;
+  while (coded_show_frame_count < num_frames) {
+    const GroupOfPicture group_of_picture =
+        simple_encode.ObserveGroupOfPicture();
+    const std::vector<EncodeFrameInfo> &encode_frame_list =
+        group_of_picture.encode_frame_list;
+    EXPECT_EQ(encode_frame_list.size(),
+              ref_gop_coded_frame_count_list[gop_count]);
+    EXPECT_EQ(group_of_picture.show_frame_count,
+              ref_gop_show_frame_count_list[gop_count]);
+    for (size_t group_index = 0; group_index < encode_frame_list.size();
+         ++group_index) {
+      EncodeFrameResult encode_frame_result;
+      simple_encode.EncodeFrame(&encode_frame_result);
+    }
+    coded_show_frame_count += group_of_picture.show_frame_count;
+    ++gop_count;
+  }
+  EXPECT_EQ(gop_count, 4);
+  simple_encode.EndEncode();
+}
+
 TEST(SimpleEncode, GetEncodeFrameInfo) {
   // Makes sure that the encode_frame_info obtained from GetEncodeFrameInfo()
   // matches the counterpart in encode_frame_result obtained from EncodeFrame()
diff --git a/vp9/simple_encode.cc b/vp9/simple_encode.cc
index f52c18015..9d62d26e1 100644
--- a/vp9/simple_encode.cc
+++ b/vp9/simple_encode.cc
@@ -806,8 +806,36 @@ std::vector<std::vector<double>> SimpleEncode::ObserveFirstPassStats() {
   return output_stats;
 }
 
-void SimpleEncode::SetExternalGroupOfPicture(std::vector<int> gop_map) {
+void SimpleEncode::SetExternalGroupOfPicturesMap(std::vector<int> gop_map) {
   gop_map_ = gop_map;
+  // The following will check and modify gop_map_ to make sure the
+  // gop_map_ satisfies the constraints.
+  // 1) Each key frame position should be at the start of a gop.
+  // 2) The last gop should not use an alt ref.
+  assert(gop_map_.size() == key_frame_map_.size());
+  int last_gop_start = 0;
+  for (int i = 0; static_cast<size_t>(i) < gop_map_.size(); ++i) {
+    if (key_frame_map_[i] == 1 && gop_map_[i] == 0) {
+      fprintf(stderr, "Add an extra gop start at show_idx %d\n", i);
+      // Insert a gop start at key frame location.
+      gop_map_[i] |= kGopMapFlagStart;
+      gop_map_[i] |= kGopMapFlagUseAltRef;
+    }
+    if (gop_map_[i] & kGopMapFlagStart) {
+      last_gop_start = i;
+    }
+  }
+  if (gop_map_[last_gop_start] & kGopMapFlagUseAltRef) {
+    fprintf(stderr,
+            "Last group of pictures starting at show_idx %d shouldn't use alt "
+            "ref\n",
+            last_gop_start);
+    gop_map_[last_gop_start] &= ~kGopMapFlagUseAltRef;
+  }
+}
+
+std::vector<int> SimpleEncode::ObserveExternalGroupOfPicturesMap() {
+  return gop_map_;
 }
 
 template <typename T>
@@ -867,6 +895,8 @@ void SimpleEncode::StartEncode() {
 
   UpdateKeyFrameGroup(show_frame_count_);
 
+  const GOP_COMMAND gop_command = GetGopCommand(gop_map_, show_frame_count_);
+  encode_command_set_gop_command(&impl_ptr_->cpi->encode_command, gop_command);
   UpdateGroupOfPicture(impl_ptr_->cpi, frame_coding_index_, ref_frame_info_,
                        &group_of_picture_);
   rewind(in_file_);
diff --git a/vp9/simple_encode.h b/vp9/simple_encode.h
index 80ecafdab..d7c9dfa14 100644
--- a/vp9/simple_encode.h
+++ b/vp9/simple_encode.h
@@ -322,7 +322,7 @@ class SimpleEncode {
   std::vector<int> ObserveKeyFrameMap() const;
 
   // Sets group of pictures map for coding the entire video.
-  // Each entry in the gop_map is corresponding to a show frame in the video.
+  // Each entry in the gop_map corresponds to a show frame in the video.
   // Therefore, the size of gop_map should equal to the number of show frames in
   // the entire video.
   // If a given entry's kGopMapFlagStart is set, it means this is the start of a
@@ -331,7 +331,16 @@ class SimpleEncode {
   // If a given entry is zero, it means it's in the middle of a gop.
   // This function should be called only once after ComputeFirstPassStats(),
   // before StartEncode().
-  void SetExternalGroupOfPicture(std::vector<int> gop_map);
+  // This API will check and modify the gop_map to satisfy the following
+  // constraints.
+  // 1) Each key frame position should be at the start of a gop.
+  // 2) The last gop should not use an alt ref.
+  void SetExternalGroupOfPicturesMap(std::vector<int> gop_map);
+
+  // Observe the group of pictures map set through
+  // SetExternalGroupOfPicturesMap(). This function should be called after
+  // SetExternalGroupOfPicturesMap().
+  std::vector<int> ObserveExternalGroupOfPicturesMap();
 
   // Initializes the encoder for actual encoding.
   // This function should be called after ComputeFirstPassStats().