Properly packetize frame bit-streams in ducky_encode

Make the bit-stream associated with the ARFs and a display frame
fit in one packet.

BUG=b/240265151

Change-Id: I2e7f2704343d25c8296ccd3de75bc77b9dcbeaa7

av1/ducky_encode.cc

diff --git a/av1/ducky_encode.cc b/av1/ducky_encode.cc
index 70de54c..6cbc91a 100644
--- a/av1/ducky_encode.cc
+++ b/av1/ducky_encode.cc
@@ -429,7 +429,9 @@
     const GopStructList &gop_list) {
   std::vector<TplGopStats> tpl_gop_stats_list;
   AV1_PRIMARY *ppi = impl_ptr_->enc_resource.ppi;
+  const VideoInfo &video_info = impl_ptr_->video_info;
   write_temp_delimiter_ = true;
+  AllocateBitstreamBuffer(video_info);
 
   // Go through each gop and encode each frame in the gop
   for (size_t i = 0; i < gop_list.size(); ++i) {
@@ -444,6 +446,7 @@
                                                   { 128, -1 } };
       (void)frame;
       EncodeFrame(frame_decision);
+      if (ppi->cpi->common.show_frame) pending_ctx_size_ = 0;
       write_temp_delimiter_ = ppi->cpi->common.show_frame;
     }
     tpl_gop_stats = ObtainTplStats(gop_struct);
@@ -460,7 +463,11 @@
     const GopEncodeInfoList &gop_encode_info_list) {
   AV1_PRIMARY *ppi = impl_ptr_->enc_resource.ppi;
   std::vector<EncodeFrameResult> encoded_frame_list;
+  const VideoInfo &video_info = impl_ptr_->video_info;
+
   write_temp_delimiter_ = true;
+  AllocateBitstreamBuffer(video_info);
+
   // Go through each gop and encode each frame in the gop
   for (size_t i = 0; i < gop_list.size(); ++i) {
     const aom::GopStruct &gop_struct = gop_list[i];
@@ -471,7 +478,15 @@
       aom::EncodeFrameDecision frame_decision = { aom::EncodeFrameMode::kQindex,
                                                   aom::EncodeGopMode::kGopRcl,
                                                   frame_param };
-      encoded_frame_list.push_back(EncodeFrame(frame_decision));
+      EncodeFrame(frame_decision);
+      if (ppi->cpi->common.show_frame) {
+        bitstream_buf_.resize(pending_ctx_size_);
+        EncodeFrameResult encode_frame_result = {};
+        encode_frame_result.bitstream_buf = bitstream_buf_;
+        encoded_frame_list.push_back(encode_frame_result);
+
+        AllocateBitstreamBuffer(video_info);
+      }
       write_temp_delimiter_ = ppi->cpi->common.show_frame;
     }
   }
@@ -482,11 +497,7 @@
 EncodeFrameResult DuckyEncode::EncodeFrame(
     const EncodeFrameDecision &decision) {
   EncodeFrameResult encode_frame_result = {};
-  const VideoInfo &video_info = impl_ptr_->video_info;
-  // Set bitstream_buf size to a conservatve upperbound.
-  // TODO(angiebird): Set bitstream_buf's size optimally.
-  encode_frame_result.bitstream_buf = std::vector<uint8_t>(
-      video_info.frame_width * video_info.frame_height * 3 * 8, 0);
+  encode_frame_result.bitstream_buf = bitstream_buf_;
   AV1_PRIMARY *ppi = impl_ptr_->enc_resource.ppi;
   aom_image_t *img = &impl_ptr_->enc_resource.img;
   AV1_COMP *const cpi = ppi->cpi;
@@ -508,8 +519,8 @@
   }
 
   AV1_COMP_DATA cpi_data = {};
-  cpi_data.cx_data = encode_frame_result.bitstream_buf.data();
-  cpi_data.cx_data_sz = encode_frame_result.bitstream_buf.size();
+  cpi_data.cx_data = bitstream_buf_.data() + pending_ctx_size_;
+  cpi_data.cx_data_sz = bitstream_buf_.size() - pending_ctx_size_;
   cpi_data.frame_size = 0;
   cpi_data.flush = 1;
   // ts_frame_start and ts_frame_end are not as important since we are focusing
@@ -529,7 +540,6 @@
   if (write_temp_delimiter_) WriteObu(ppi, &cpi_data);
   (void)status;
   assert(status == static_cast<int>(AOM_CODEC_OK));
-  encode_frame_result.bitstream_buf.resize(cpi_data.frame_size);
   DuckyEncodeInfoGetEncodeFrameResult(&cpi->ducky_encode_info,
                                       &encode_frame_result);
   av1_post_encode_updates(cpi, &cpi_data);
@@ -538,6 +548,8 @@
     ppi->frames_left = AOMMAX(0, ppi->frames_left - 1);
   }
 
+  pending_ctx_size_ += cpi_data.frame_size;
+
   fprintf(stderr, "frame %d, qp = %d, size %d, PSNR %f\n",
           encode_frame_result.global_order_idx, encode_frame_result.q_index,
           encode_frame_result.rate, encode_frame_result.psnr);
@@ -545,4 +557,11 @@
 }
 
 void DuckyEncode::EndEncode() { FreeEncoder(&impl_ptr_->enc_resource); }
+
+void DuckyEncode::AllocateBitstreamBuffer(const VideoInfo &video_info) {
+  pending_ctx_size_ = 0;
+  // TODO(angiebird): Set bitstream_buf size to a conservatve upperbound.
+  bitstream_buf_.assign(
+      video_info.frame_width * video_info.frame_height * 3 * 8, 0);
+}
 }  // namespace aom

av1/ducky_encode.h

diff --git a/av1/ducky_encode.h b/av1/ducky_encode.h
index 1acb31f..29dc29e 100644
--- a/av1/ducky_encode.h
+++ b/av1/ducky_encode.h
@@ -83,11 +83,14 @@
       const GopEncodeInfoList &gop_encode_info_list);
   EncodeFrameResult EncodeFrame(const EncodeFrameDecision &decision);
   void EndEncode();
+  void AllocateBitstreamBuffer(const VideoInfo &video_info);
 
  private:
   class EncodeImpl;
   std::unique_ptr<EncodeImpl> impl_ptr_;
   bool write_temp_delimiter_;
+  std::vector<uint8_t> bitstream_buf_;
+  int pending_ctx_size_;
 };
 }  // namespace aom
 

test/ducky_encode_test.cc

diff --git a/test/ducky_encode_test.cc b/test/ducky_encode_test.cc
index 4b2d3ed..ce64253 100644
--- a/test/ducky_encode_test.cc
+++ b/test/ducky_encode_test.cc
@@ -65,6 +65,7 @@
                                    aom::EncodeGopMode::kNone,
                                    {} };
   for (int i = 0; i < coding_frame_count; ++i) {
+    ducky_encode.AllocateBitstreamBuffer(video_info);
     EncodeFrameResult encode_frame_result = ducky_encode.EncodeFrame(decision);
   }
   ducky_encode.EndEncode();
@@ -91,6 +92,7 @@
                                    aom::EncodeGopMode::kNone,
                                    { q_index, -1 } };
   for (int i = 0; i < coding_frame_count; ++i) {
+    ducky_encode.AllocateBitstreamBuffer(video_info);
     EncodeFrameResult encode_frame_result = ducky_encode.EncodeFrame(decision);
     // TODO(angiebird): Check why distortion is not zero when q_index = 0
     EXPECT_EQ(encode_frame_result.dist, 0);