twopass_encoder: Flush encoder.

Change-Id: Ib04a545b43f3a36c6df7eaf3ed2d802cf41119bb

1 files changed, 33 insertions, 17 deletions

diff --git a/examples/twopass_encoder.c b/examples/twopass_encoder.c
index 74291f920..76d5a2816 100644
--- a/examples/twopass_encoder.c
+++ b/examples/twopass_encoder.c
@@ -66,13 +66,14 @@ void usage_exit() {
   exit(EXIT_FAILURE);
 }
 
-static void get_frame_stats(vpx_codec_ctx_t *ctx,
-                            const vpx_image_t *img,
-                            vpx_codec_pts_t pts,
-                            unsigned int duration,
-                            vpx_enc_frame_flags_t flags,
-                            unsigned int deadline,
-                            vpx_fixed_buf_t *stats) {
+static int get_frame_stats(vpx_codec_ctx_t *ctx,
+                           const vpx_image_t *img,
+                           vpx_codec_pts_t pts,
+                           unsigned int duration,
+                           vpx_enc_frame_flags_t flags,
+                           unsigned int deadline,
+                           vpx_fixed_buf_t *stats) {
+  int got_pkts = 0;
   vpx_codec_iter_t iter = NULL;
   const vpx_codec_cx_pkt_t *pkt = NULL;
   const vpx_codec_err_t res = vpx_codec_encode(ctx, img, pts, duration, flags,
@@ -81,6 +82,8 @@ static void get_frame_stats(vpx_codec_ctx_t *ctx,
     die_codec(ctx, "Failed to get frame stats.");
 
   while ((pkt = vpx_codec_get_cx_data(ctx, &iter)) != NULL) {
+    got_pkts = 1;
+
     if (pkt->kind == VPX_CODEC_STATS_PKT) {
       const uint8_t *const pkt_buf = pkt->data.twopass_stats.buf;
       const size_t pkt_size = pkt->data.twopass_stats.sz;
@@ -89,15 +92,18 @@ static void get_frame_stats(vpx_codec_ctx_t *ctx,
       stats->sz += pkt_size;
     }
   }
+
+  return got_pkts;
 }
 
-static void encode_frame(vpx_codec_ctx_t *ctx,
-                         const vpx_image_t *img,
-                         vpx_codec_pts_t pts,
-                         unsigned int duration,
-                         vpx_enc_frame_flags_t flags,
-                         unsigned int deadline,
-                         VpxVideoWriter *writer) {
+static int encode_frame(vpx_codec_ctx_t *ctx,
+                        const vpx_image_t *img,
+                        vpx_codec_pts_t pts,
+                        unsigned int duration,
+                        vpx_enc_frame_flags_t flags,
+                        unsigned int deadline,
+                        VpxVideoWriter *writer) {
+  int got_pkts = 0;
   vpx_codec_iter_t iter = NULL;
   const vpx_codec_cx_pkt_t *pkt = NULL;
   const vpx_codec_err_t res = vpx_codec_encode(ctx, img, pts, duration, flags,
@@ -106,6 +112,7 @@ static void encode_frame(vpx_codec_ctx_t *ctx,
     die_codec(ctx, "Failed to encode frame.");
 
   while ((pkt = vpx_codec_get_cx_data(ctx, &iter)) != NULL) {
+    got_pkts = 1;
     if (pkt->kind == VPX_CODEC_CX_FRAME_PKT) {
       const int keyframe = (pkt->data.frame.flags & VPX_FRAME_IS_KEY) != 0;
 
@@ -117,12 +124,14 @@ static void encode_frame(vpx_codec_ctx_t *ctx,
       fflush(stdout);
     }
   }
+
+  return got_pkts;
 }
 
 static vpx_fixed_buf_t pass0(vpx_image_t *raw,
                              FILE *infile,
                              const VpxInterface *encoder,
-                             vpx_codec_enc_cfg_t *cfg) {
+                             const vpx_codec_enc_cfg_t *cfg) {
   vpx_codec_ctx_t codec;
   int frame_count = 0;
   vpx_fixed_buf_t stats = {NULL, 0};
@@ -130,13 +139,16 @@ static vpx_fixed_buf_t pass0(vpx_image_t *raw,
   if (vpx_codec_enc_init(&codec, encoder->codec_interface(), cfg, 0))
     die_codec(&codec, "Failed to initialize encoder");
 
+  // Calculate frame statistics.
   while (vpx_img_read(raw, infile)) {
     ++frame_count;
     get_frame_stats(&codec, raw, frame_count, 1, 0, VPX_DL_BEST_QUALITY,
                     &stats);
   }
 
-  get_frame_stats(&codec, NULL, frame_count, 1, 0, VPX_DL_BEST_QUALITY, &stats);
+  // Flush encoder.
+  while (get_frame_stats(&codec, NULL, frame_count, 1, 0,
+                         VPX_DL_BEST_QUALITY, &stats)) {}
 
   printf("Pass 0 complete. Processed %d frames.\n", frame_count);
   if (vpx_codec_destroy(&codec))
@@ -149,7 +161,7 @@ static void pass1(vpx_image_t *raw,
                   FILE *infile,
                   const char *outfile_name,
                   const VpxInterface *encoder,
-                  vpx_codec_enc_cfg_t *cfg) {
+                  const vpx_codec_enc_cfg_t *cfg) {
   VpxVideoInfo info = {
     encoder->fourcc,
     cfg->g_w,
@@ -167,11 +179,15 @@ static void pass1(vpx_image_t *raw,
   if (vpx_codec_enc_init(&codec, encoder->codec_interface(), cfg, 0))
     die_codec(&codec, "Failed to initialize encoder");
 
+  // Encode frames.
   while (vpx_img_read(raw, infile)) {
     ++frame_count;
     encode_frame(&codec, raw, frame_count, 1, 0, VPX_DL_BEST_QUALITY, writer);
   }
 
+  // Flush encoder.
+  while (encode_frame(&codec, NULL, -1, 1, 0, VPX_DL_BEST_QUALITY, writer)) {}
+
   printf("\n");
 
   if (vpx_codec_destroy(&codec))