src/third_party/libvpx/source/libvpx/examples/twopass_encoder.c

   1 /*
   2  *  Copyright (c) 2010 The WebM project authors. All Rights Reserved.
   3  *
   4  *  Use of this source code is governed by a BSD-style license
   5  *  that can be found in the LICENSE file in the root of the source
   6  *  tree. An additional intellectual property rights grant can be found
   7  *  in the file PATENTS.  All contributing project authors may
   8  *  be found in the AUTHORS file in the root of the source tree.
   9  */
  10
  11 // Two Pass Encoder
  12 // ================
  13 //
  14 // This is an example of a two pass encoder loop. It takes an input file in
  15 // YV12 format, passes it through the encoder twice, and writes the compressed
  16 // frames to disk in IVF format. It builds upon the simple_encoder example.
  17 //
  18 // Twopass Variables
  19 // -----------------
  20 // Twopass mode needs to track the current pass number and the buffer of
  21 // statistics packets.
  22 //
  23 // Updating The Configuration
  24 // ---------------------------------
  25 // In two pass mode, the configuration has to be updated on each pass. The
  26 // statistics buffer is passed on the last pass.
  27 //
  28 // Encoding A Frame
  29 // ----------------
  30 // Encoding a frame in two pass mode is identical to the simple encoder
  31 // example, except the deadline is set to VPX_DL_BEST_QUALITY to get the
  32 // best quality possible. VPX_DL_GOOD_QUALITY could also be used.
  33 //
  34 //
  35 // Processing Statistics Packets
  36 // -----------------------------
  37 // Each packet of type `VPX_CODEC_CX_FRAME_PKT` contains the encoded data
  38 // for this frame. We write a IVF frame header, followed by the raw data.
  39 //
  40 //
  41 // Pass Progress Reporting
  42 // -----------------------------
  43 // It's sometimes helpful to see when each pass completes.
  44 //
  45 //
  46 // Clean-up
  47 // -----------------------------
  48 // Destruction of the encoder instance must be done on each pass. The
  49 // raw image should be destroyed at the end as usual.
  50
  51 #include <stdio.h>
  52 #include <stdlib.h>
  53 #include <string.h>
  54
  55 #define VPX_CODEC_DISABLE_COMPAT 1
  56 #include "vpx/vpx_encoder.h"
  57
  58 #include "./tools_common.h"
  59 #include "./video_writer.h"
  60
  61 static const char *exec_name;
  62
  63 void usage_exit() {
  64   fprintf(stderr, "Usage: %s <codec> <width> <height> <infile> <outfile>\n",
  65           exec_name);
  66   exit(EXIT_FAILURE);
  67 }
  68
  69 static void get_frame_stats(vpx_codec_ctx_t *ctx,
  70                             const vpx_image_t *img,
  71                             vpx_codec_pts_t pts,
  72                             uint64_t duration,
  73                             vpx_enc_frame_flags_t flags,
  74                             uint64_t deadline,
  75                             vpx_fixed_buf_t *stats) {
  76   vpx_codec_iter_t iter = NULL;
  77   const vpx_codec_cx_pkt_t *pkt = NULL;
  78   const vpx_codec_err_t res = vpx_codec_encode(ctx, img, pts, duration, flags,
  79                                                deadline);
  80   if (res != VPX_CODEC_OK)
  81     die_codec(ctx, "Failed to get frame stats.");
  82
  83   while ((pkt = vpx_codec_get_cx_data(ctx, &iter)) != NULL) {
  84     if (pkt->kind == VPX_CODEC_STATS_PKT) {
  85       const uint8_t *const pkt_buf = pkt->data.twopass_stats.buf;
  86       const size_t pkt_size = pkt->data.twopass_stats.sz;
  87       stats->buf = realloc(stats->buf, stats->sz + pkt_size);
  88       memcpy((uint8_t *)stats->buf + stats->sz, pkt_buf, pkt_size);
  89       stats->sz += pkt_size;
  90     }
  91   }
  92 }
  93
  94 static void encode_frame(vpx_codec_ctx_t *ctx,
  95                          const vpx_image_t *img,
  96                          vpx_codec_pts_t pts,
  97                          uint64_t duration,
  98                          vpx_enc_frame_flags_t flags,
  99                          uint64_t deadline,
 100                          VpxVideoWriter *writer) {
 101   vpx_codec_iter_t iter = NULL;
 102   const vpx_codec_cx_pkt_t *pkt = NULL;
 103   const vpx_codec_err_t res = vpx_codec_encode(ctx, img, pts, duration, flags,
 104                                                deadline);
 105   if (res != VPX_CODEC_OK)
 106     die_codec(ctx, "Failed to encode frame.");
 107
 108   while ((pkt = vpx_codec_get_cx_data(ctx, &iter)) != NULL) {
 109     if (pkt->kind == VPX_CODEC_CX_FRAME_PKT) {
 110       const int keyframe = (pkt->data.frame.flags & VPX_FRAME_IS_KEY) != 0;
 111
 112       if (!vpx_video_writer_write_frame(writer, pkt->data.frame.buf,
 113                                                 pkt->data.frame.sz,
 114                                                 pkt->data.frame.pts))
 115         die_codec(ctx, "Failed to write compressed frame.");
 116       printf(keyframe ? "K" : ".");
 117       fflush(stdout);
 118     }
 119   }
 120 }
 121
 122 int main(int argc, char **argv) {
 123   FILE *infile = NULL;
 124   VpxVideoWriter *writer = NULL;
 125   vpx_codec_ctx_t codec;
 126   vpx_codec_enc_cfg_t cfg;
 127   vpx_image_t raw;
 128   vpx_codec_err_t res;
 129   vpx_fixed_buf_t stats = {0};
 130   VpxVideoInfo info = {0};
 131   const VpxInterface *encoder = NULL;
 132   int pass;
 133   const int fps = 30;        // TODO(dkovalev) add command line argument
 134   const int bitrate = 200;   // kbit/s TODO(dkovalev) add command line argument
 135   const char *const codec_arg = argv[1];
 136   const char *const width_arg = argv[2];
 137   const char *const height_arg = argv[3];
 138   const char *const infile_arg = argv[4];
 139   const char *const outfile_arg = argv[5];
 140   exec_name = argv[0];
 141
 142   if (argc != 6)
 143     die("Invalid number of arguments.");
 144
 145   encoder = get_vpx_encoder_by_name(codec_arg);
 146   if (!encoder)
 147     die("Unsupported codec.");
 148
 149   info.codec_fourcc = encoder->fourcc;
 150   info.time_base.numerator = 1;
 151   info.time_base.denominator = fps;
 152   info.frame_width = strtol(width_arg, NULL, 0);
 153   info.frame_height = strtol(height_arg, NULL, 0);
 154
 155   if (info.frame_width <= 0 ||
 156       info.frame_height <= 0 ||
 157       (info.frame_width % 2) != 0 ||
 158       (info.frame_height % 2) != 0) {
 159     die("Invalid frame size: %dx%d", info.frame_width, info.frame_height);
 160   }
 161
 162   if (!vpx_img_alloc(&raw, VPX_IMG_FMT_I420, info.frame_width,
 163                                              info.frame_height, 1)) {
 164     die("Failed to allocate image", info.frame_width, info.frame_height);
 165   }
 166
 167   writer = vpx_video_writer_open(outfile_arg, kContainerIVF, &info);
 168   if (!writer)
 169     die("Failed to open %s for writing", outfile_arg);
 170
 171   printf("Using %s\n", vpx_codec_iface_name(encoder->interface()));
 172
 173   res = vpx_codec_enc_config_default(encoder->interface(), &cfg, 0);
 174   if (res)
 175     die_codec(&codec, "Failed to get default codec config.");
 176
 177   cfg.g_w = info.frame_width;
 178   cfg.g_h = info.frame_height;
 179   cfg.g_timebase.num = info.time_base.numerator;
 180   cfg.g_timebase.den = info.time_base.denominator;
 181   cfg.rc_target_bitrate = bitrate;
 182
 183   for (pass = 0; pass < 2; ++pass) {
 184     int frame_count = 0;
 185
 186     if (pass == 0) {
 187       cfg.g_pass = VPX_RC_FIRST_PASS;
 188     } else {
 189       cfg.g_pass = VPX_RC_LAST_PASS;
 190       cfg.rc_twopass_stats_in = stats;
 191     }
 192
 193     if (!(infile = fopen(infile_arg, "rb")))
 194       die("Failed to open %s for reading", infile_arg);
 195
 196     if (vpx_codec_enc_init(&codec, encoder->interface(), &cfg, 0))
 197       die_codec(&codec, "Failed to initialize encoder");
 198
 199     while (vpx_img_read(&raw, infile)) {
 200       ++frame_count;
 201
 202       if (pass == 0) {
 203         get_frame_stats(&codec, &raw, frame_count, 1, 0, VPX_DL_BEST_QUALITY,
 204                         &stats);
 205       } else {
 206         encode_frame(&codec, &raw, frame_count, 1, 0, VPX_DL_BEST_QUALITY,
 207                      writer);
 208       }
 209     }
 210
 211     if (pass == 0) {
 212       get_frame_stats(&codec, NULL, frame_count, 1, 0, VPX_DL_BEST_QUALITY,
 213                       &stats);
 214     } else {
 215       printf("\n");
 216     }
 217
 218     fclose(infile);
 219     printf("Pass %d complete. Processed %d frames.\n", pass + 1, frame_count);
 220     if (vpx_codec_destroy(&codec))
 221       die_codec(&codec, "Failed to destroy codec.");
 222   }
 223
 224   vpx_img_free(&raw);
 225   free(stats.buf);
 226
 227   vpx_video_writer_close(writer);
 228
 229   return EXIT_SUCCESS;
 230 }