Add functions to compute/observe key frame map

Change-Id: I2fc0efb2ac35e64af3350bddaa802a206d1aa13c

5 files changed, 104 insertions, 3 deletions

diff --git a/test/simple_encode_test.cc b/test/simple_encode_test.cc
index 45dcad050..cbcae46be 100644
--- a/test/simple_encode_test.cc
+++ b/test/simple_encode_test.cc
@@ -106,6 +106,34 @@ TEST(SimpleEncode, EncodeFrame) {
   simple_encode.EndEncode();
 }
 
+TEST(SimpleEncode, ObserveKeyFrameMap) {
+  SimpleEncode simple_encode(w, h, frame_rate_num, frame_rate_den,
+                             target_bitrate, num_frames, infile_path);
+  simple_encode.ComputeFirstPassStats();
+  std::vector<int> key_frame_map = simple_encode.ObserveKeyFrameMap();
+  EXPECT_EQ(key_frame_map.size(), static_cast<size_t>(num_frames));
+  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;
+    for (size_t group_index = 0; group_index < encode_frame_list.size();
+         ++group_index) {
+      EncodeFrameResult encode_frame_result;
+      simple_encode.EncodeFrame(&encode_frame_result);
+      if (encode_frame_result.frame_type == kFrameTypeKey) {
+        EXPECT_EQ(key_frame_map[encode_frame_result.show_idx], 1);
+      } else {
+        EXPECT_EQ(key_frame_map[encode_frame_result.show_idx], 0);
+      }
+    }
+    coded_show_frame_count += group_of_picture.show_frame_count;
+  }
+  simple_encode.EndEncode();
+}
+
 TEST(SimpleEncode, EncodeFrameWithQuantizeIndex) {
   SimpleEncode simple_encode(w, h, frame_rate_num, frame_rate_den,
                              target_bitrate, num_frames, infile_path);
diff --git a/vp9/encoder/vp9_firstpass.c b/vp9/encoder/vp9_firstpass.c
index 4bb58ecf4..45b003e1e 100644
--- a/vp9/encoder/vp9_firstpass.c
+++ b/vp9/encoder/vp9_firstpass.c
@@ -3796,6 +3796,30 @@ int vp9_get_coding_frame_num(const VP9EncoderConfig *oxcf,
   }
   return coding_frame_num;
 }
+
+void vp9_get_key_frame_map(const VP9EncoderConfig *oxcf,
+                           const FRAME_INFO *frame_info,
+                           const FIRST_PASS_INFO *first_pass_info,
+                           int *key_frame_map) {
+  int show_idx = 0;
+  RATE_CONTROL rc;
+  vp9_rc_init(oxcf, 1, &rc);
+
+  // key_frame_map points to an int array with size equal to
+  // first_pass_info->num_frames, which is also the number of show frames in the
+  // video.
+  memset(key_frame_map, 0,
+         sizeof(*key_frame_map) * first_pass_info->num_frames);
+  while (show_idx < first_pass_info->num_frames) {
+    int key_frame_group_size;
+    key_frame_map[show_idx] = 1;
+    key_frame_group_size = vp9_get_frames_to_next_key(
+        oxcf, frame_info, first_pass_info, show_idx, rc.min_gf_interval);
+    assert(key_frame_group_size > 0);
+    show_idx += key_frame_group_size;
+  }
+  assert(show_idx == first_pass_info->num_frames);
+}
 #endif  // CONFIG_RATE_CTRL
 
 FIRSTPASS_STATS vp9_get_frame_stats(const TWO_PASS *twopass) {
diff --git a/vp9/encoder/vp9_firstpass.h b/vp9/encoder/vp9_firstpass.h
index 51ffd0222..b1047eab2 100644
--- a/vp9/encoder/vp9_firstpass.h
+++ b/vp9/encoder/vp9_firstpass.h
@@ -290,6 +290,16 @@ int vp9_get_coding_frame_num(const struct VP9EncoderConfig *oxcf,
                              const FRAME_INFO *frame_info,
                              const FIRST_PASS_INFO *first_pass_info,
                              int multi_layer_arf, int allow_alt_ref);
+
+/*!\brief Compute a key frame binary map indicates whether key frames appear
+ * in the corresponding positions. The passed in key_frame_map must point to an
+ * integer array with length equal to first_pass_info->num_frames, which is the
+ * number of show frames in the video.
+ */
+void vp9_get_key_frame_map(const struct VP9EncoderConfig *oxcf,
+                           const FRAME_INFO *frame_info,
+                           const FIRST_PASS_INFO *first_pass_info,
+                           int *key_frame_map);
 #endif  // CONFIG_RATE_CTRL
 
 FIRSTPASS_STATS vp9_get_frame_stats(const TWO_PASS *twopass);
diff --git a/vp9/simple_encode.cc b/vp9/simple_encode.cc
index 7e6ecf8ea..f52c18015 100644
--- a/vp9/simple_encode.cc
+++ b/vp9/simple_encode.cc
@@ -779,6 +779,9 @@ void SimpleEncode::ComputeFirstPassStats() {
   free_encoder(cpi);
   rewind(in_file_);
   vpx_img_free(&img);
+
+  // Generate key_frame_map based on impl_ptr_->first_pass_stats.
+  key_frame_map_ = ComputeKeyFrameMap();
 }
 
 std::vector<std::vector<double>> SimpleEncode::ObserveFirstPassStats() {
@@ -1065,6 +1068,28 @@ int SimpleEncode::GetCodingFrameNum() const {
                                   multi_layer_arf, allow_alt_ref);
 }
 
+std::vector<int> SimpleEncode::ComputeKeyFrameMap() const {
+  assert(impl_ptr_->first_pass_stats.size() == num_frames_);
+  vpx_rational_t frame_rate =
+      make_vpx_rational(frame_rate_num_, frame_rate_den_);
+  const VP9EncoderConfig oxcf =
+      vp9_get_encoder_config(frame_width_, frame_height_, frame_rate,
+                             target_bitrate_, VPX_RC_LAST_PASS);
+  FRAME_INFO frame_info = vp9_get_frame_info(&oxcf);
+  FIRST_PASS_INFO first_pass_info;
+  fps_init_first_pass_info(&first_pass_info,
+                           GetVectorData(impl_ptr_->first_pass_stats),
+                           num_frames_);
+  std::vector<int> key_frame_map(num_frames_, 0);
+  vp9_get_key_frame_map(&oxcf, &frame_info, &first_pass_info,
+                        GetVectorData(key_frame_map));
+  return key_frame_map;
+}
+
+std::vector<int> SimpleEncode::ObserveKeyFrameMap() const {
+  return key_frame_map_;
+}
+
 uint64_t SimpleEncode::GetFramePixelCount() const {
   assert(frame_width_ % 2 == 0);
   assert(frame_height_ % 2 == 0);
diff --git a/vp9/simple_encode.h b/vp9/simple_encode.h
index d84760ac9..80ecafdab 100644
--- a/vp9/simple_encode.h
+++ b/vp9/simple_encode.h
@@ -304,8 +304,9 @@ class SimpleEncode {
   SimpleEncode(SimpleEncode &) = delete;
   SimpleEncode &operator=(const SimpleEncode &) = delete;
 
-  // Makes encoder compute the first pass stats and store it internally for
-  // future encode.
+  // Makes encoder compute the first pass stats and store it at
+  // impl_ptr_->first_pass_stats. key_frame_map_ is also computed based on the
+  // first pass stats.
   void ComputeFirstPassStats();
 
   // Outputs the first pass stats represented by a 2-D vector.
@@ -314,8 +315,14 @@ class SimpleEncode {
   // values. For details, please check FIRSTPASS_STATS in vp9_firstpass.h
   std::vector<std::vector<double>> ObserveFirstPassStats();
 
+  // Ouputs a copy of key_frame_map_, a binary vector with size equal to the
+  // number of show frames in the video. For each entry in the vector, 1
+  // indicates the position is a key frame and 0 indicates it's not a key frame.
+  // This function should be called after ComputeFirstPassStats()
+  std::vector<int> ObserveKeyFrameMap() const;
+
   // Sets group of pictures map for coding the entire video.
-  // Each entry in the gop_map corresponds to a show frame in the video.
+  // Each entry in the gop_map is corresponding 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
@@ -366,6 +373,12 @@ class SimpleEncode {
   uint64_t GetFramePixelCount() const;
 
  private:
+  // Compute the key frame locations of the video based on first pass stats.
+  // The results are returned as a binary vector with 1s indicating keyframes
+  // and 0s indicating non keyframes.
+  // It has to be called after impl_ptr_->first_pass_stats is computed.
+  std::vector<int> ComputeKeyFrameMap() const;
+
   // Updates key_frame_group_size_, reset key_frame_group_index_ and init
   // ref_frame_info_.
   void UpdateKeyFrameGroup(int key_frame_show_index);
@@ -388,6 +401,7 @@ class SimpleEncode {
   std::FILE *out_file_;
   std::unique_ptr<EncodeImpl> impl_ptr_;
 
+  std::vector<int> key_frame_map_;
   std::vector<int> gop_map_;
   GroupOfPicture group_of_picture_;