From 23b070f46ed9d411f547d6481b157efc60e2d5d8 Mon Sep 17 00:00:00 2001 From: angiebird Date: Tue, 26 May 2020 19:02:33 -0700 Subject: Add functions to compute/observe key frame map Change-Id: I2fc0efb2ac35e64af3350bddaa802a206d1aa13c --- test/simple_encode_test.cc | 28 ++++++++++++++++++++++++++++ vp9/encoder/vp9_firstpass.c | 24 ++++++++++++++++++++++++ vp9/encoder/vp9_firstpass.h | 10 ++++++++++ vp9/simple_encode.cc | 25 +++++++++++++++++++++++++ vp9/simple_encode.h | 20 +++++++++++++++++--- 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 key_frame_map = simple_encode.ObserveKeyFrameMap(); + EXPECT_EQ(key_frame_map.size(), static_cast(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 &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> SimpleEncode::ObserveFirstPassStats() { @@ -1065,6 +1068,28 @@ int SimpleEncode::GetCodingFrameNum() const { multi_layer_arf, allow_alt_ref); } +std::vector 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 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 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> 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 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 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 impl_ptr_; + std::vector key_frame_map_; std::vector gop_map_; GroupOfPicture group_of_picture_; -- cgit v1.2.1