shithub: libvpx

--- a/test/simple_encode_test.cc

+++ b/test/simple_encode_test.cc

@@ -106,6 +106,34 @@

   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);

--- a/vp9/encoder/vp9_firstpass.c

+++ b/vp9/encoder/vp9_firstpass.c

@@ -3796,6 +3796,30 @@

   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) {

--- a/vp9/encoder/vp9_firstpass.h

+++ b/vp9/encoder/vp9_firstpass.h

@@ -290,6 +290,16 @@

                              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);

--- a/vp9/simple_encode.cc

+++ b/vp9/simple_encode.cc

@@ -779,6 +779,9 @@

   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() {

@@ -1063,6 +1066,28 @@

                            num_frames_);

   return vp9_get_coding_frame_num(&oxcf, &frame_info, &first_pass_info,

                                   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 {

--- a/vp9/simple_encode.h

+++ b/vp9/simple_encode.h

@@ -304,8 +304,9 @@

   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 @@

   // 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 @@

   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 @@

   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_;