twopass_encoder.c (8186B)
1 /* 2 * Copyright (c) 2016, Alliance for Open Media. All rights reserved. 3 * 4 * This source code is subject to the terms of the BSD 2 Clause License and 5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License 6 * was not distributed with this source code in the LICENSE file, you can 7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open 8 * Media Patent License 1.0 was not distributed with this source code in the 9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent. 10 */ 11 12 // Two Pass Encoder 13 // ================ 14 // 15 // This is an example of a two pass encoder loop. It takes an input file in 16 // YV12 format, passes it through the encoder twice, and writes the compressed 17 // frames to disk in IVF format. It builds upon the simple_encoder example. 18 // 19 // Twopass Variables 20 // ----------------- 21 // Twopass mode needs to track the current pass number and the buffer of 22 // statistics packets. 23 // 24 // Updating The Configuration 25 // --------------------------------- 26 // In two pass mode, the configuration has to be updated on each pass. The 27 // statistics buffer is passed on the last pass. 28 // 29 // Encoding A Frame 30 // ---------------- 31 // Encoding a frame in two pass mode is identical to the simple encoder 32 // example. 33 // 34 // Processing Statistics Packets 35 // ----------------------------- 36 // Each packet of type `AOM_CODEC_CX_FRAME_PKT` contains the encoded data 37 // for this frame. We write a IVF frame header, followed by the raw data. 38 // 39 // 40 // Pass Progress Reporting 41 // ----------------------------- 42 // It's sometimes helpful to see when each pass completes. 43 // 44 // 45 // Clean-up 46 // ----------------------------- 47 // Destruction of the encoder instance must be done on each pass. The 48 // raw image should be destroyed at the end as usual. 49 50 #include <stdio.h> 51 #include <stdlib.h> 52 #include <string.h> 53 54 #include "aom/aom_encoder.h" 55 #include "aom/aomcx.h" 56 #include "common/tools_common.h" 57 #include "common/video_writer.h" 58 59 static const char *exec_name; 60 61 void usage_exit(void) { 62 fprintf(stderr, 63 "Usage: %s <codec> <width> <height> <infile> <outfile> " 64 "<limit(optional)>\n", 65 exec_name); 66 exit(EXIT_FAILURE); 67 } 68 69 static int get_frame_stats(aom_codec_ctx_t *ctx, const aom_image_t *img, 70 aom_codec_pts_t pts, unsigned int duration, 71 aom_enc_frame_flags_t flags, 72 aom_fixed_buf_t *stats) { 73 int got_pkts = 0; 74 aom_codec_iter_t iter = NULL; 75 const aom_codec_cx_pkt_t *pkt = NULL; 76 const aom_codec_err_t res = aom_codec_encode(ctx, img, pts, duration, flags); 77 if (res != AOM_CODEC_OK) die_codec(ctx, "Failed to get frame stats."); 78 79 while ((pkt = aom_codec_get_cx_data(ctx, &iter)) != NULL) { 80 got_pkts = 1; 81 82 if (pkt->kind == AOM_CODEC_STATS_PKT) { 83 const uint8_t *const pkt_buf = pkt->data.twopass_stats.buf; 84 const size_t pkt_size = pkt->data.twopass_stats.sz; 85 stats->buf = realloc(stats->buf, stats->sz + pkt_size); 86 if (!stats->buf) die("Failed to allocate frame stats buffer."); 87 memcpy((uint8_t *)stats->buf + stats->sz, pkt_buf, pkt_size); 88 stats->sz += pkt_size; 89 } 90 } 91 92 return got_pkts; 93 } 94 95 static int encode_frame(aom_codec_ctx_t *ctx, const aom_image_t *img, 96 aom_codec_pts_t pts, unsigned int duration, 97 aom_enc_frame_flags_t flags, AvxVideoWriter *writer) { 98 int got_pkts = 0; 99 aom_codec_iter_t iter = NULL; 100 const aom_codec_cx_pkt_t *pkt = NULL; 101 const aom_codec_err_t res = aom_codec_encode(ctx, img, pts, duration, flags); 102 if (res != AOM_CODEC_OK) die_codec(ctx, "Failed to encode frame."); 103 104 while ((pkt = aom_codec_get_cx_data(ctx, &iter)) != NULL) { 105 got_pkts = 1; 106 if (pkt->kind == AOM_CODEC_CX_FRAME_PKT) { 107 const int keyframe = (pkt->data.frame.flags & AOM_FRAME_IS_KEY) != 0; 108 109 if (!aom_video_writer_write_frame(writer, pkt->data.frame.buf, 110 pkt->data.frame.sz, 111 pkt->data.frame.pts)) 112 die_codec(ctx, "Failed to write compressed frame."); 113 printf(keyframe ? "K" : "."); 114 fflush(stdout); 115 } 116 } 117 118 return got_pkts; 119 } 120 121 static aom_fixed_buf_t pass0(aom_image_t *raw, FILE *infile, 122 aom_codec_iface_t *encoder, 123 const aom_codec_enc_cfg_t *cfg, int limit) { 124 aom_codec_ctx_t codec; 125 int frame_count = 0; 126 aom_fixed_buf_t stats = { NULL, 0 }; 127 128 if (aom_codec_enc_init(&codec, encoder, cfg, 0)) 129 die("Failed to initialize encoder"); 130 131 // Calculate frame statistics. 132 while (aom_img_read(raw, infile) && frame_count < limit) { 133 ++frame_count; 134 get_frame_stats(&codec, raw, frame_count, 1, 0, &stats); 135 } 136 137 // Flush encoder. 138 while (get_frame_stats(&codec, NULL, frame_count, 1, 0, &stats)) { 139 } 140 141 printf("Pass 0 complete. Processed %d frames.\n", frame_count); 142 if (aom_codec_destroy(&codec)) die_codec(&codec, "Failed to destroy codec."); 143 144 return stats; 145 } 146 147 static void pass1(aom_image_t *raw, FILE *infile, const char *outfile_name, 148 aom_codec_iface_t *encoder, const aom_codec_enc_cfg_t *cfg, 149 int limit) { 150 AvxVideoInfo info = { get_fourcc_by_aom_encoder(encoder), 151 cfg->g_w, 152 cfg->g_h, 153 { cfg->g_timebase.num, cfg->g_timebase.den }, 154 0 }; 155 AvxVideoWriter *writer = NULL; 156 aom_codec_ctx_t codec; 157 int frame_count = 0; 158 159 writer = aom_video_writer_open(outfile_name, kContainerIVF, &info); 160 if (!writer) die("Failed to open %s for writing", outfile_name); 161 162 if (aom_codec_enc_init(&codec, encoder, cfg, 0)) 163 die("Failed to initialize encoder"); 164 165 if (aom_codec_control(&codec, AOME_SET_CPUUSED, 2)) 166 die_codec(&codec, "Failed to set cpu-used"); 167 168 // Encode frames. 169 while (aom_img_read(raw, infile) && frame_count < limit) { 170 ++frame_count; 171 encode_frame(&codec, raw, frame_count, 1, 0, writer); 172 } 173 174 // Flush encoder. 175 while (encode_frame(&codec, NULL, -1, 1, 0, writer)) { 176 } 177 178 printf("\n"); 179 180 if (aom_codec_destroy(&codec)) die_codec(&codec, "Failed to destroy codec."); 181 182 aom_video_writer_close(writer); 183 184 printf("Pass 1 complete. Processed %d frames.\n", frame_count); 185 } 186 187 int main(int argc, char **argv) { 188 FILE *infile = NULL; 189 int w, h; 190 aom_codec_ctx_t codec; 191 aom_codec_enc_cfg_t cfg; 192 aom_image_t raw; 193 aom_codec_err_t res; 194 aom_fixed_buf_t stats; 195 196 const int fps = 30; // TODO(dkovalev) add command line argument 197 const int bitrate = 200; // kbit/s TODO(dkovalev) add command line argument 198 const char *const codec_arg = argv[1]; 199 const char *const width_arg = argv[2]; 200 const char *const height_arg = argv[3]; 201 const char *const infile_arg = argv[4]; 202 const char *const outfile_arg = argv[5]; 203 int limit = 0; 204 exec_name = argv[0]; 205 206 if (argc < 6) die("Invalid number of arguments"); 207 208 if (argc > 6) limit = (int)strtol(argv[6], NULL, 0); 209 210 if (limit == 0) limit = 100; 211 212 aom_codec_iface_t *encoder = get_aom_encoder_by_short_name(codec_arg); 213 if (!encoder) die("Unsupported codec."); 214 215 w = (int)strtol(width_arg, NULL, 0); 216 h = (int)strtol(height_arg, NULL, 0); 217 218 if (w <= 0 || h <= 0 || (w % 2) != 0 || (h % 2) != 0) 219 die("Invalid frame size: %dx%d", w, h); 220 221 if (!aom_img_alloc(&raw, AOM_IMG_FMT_I420, w, h, 1)) 222 die("Failed to allocate image (%dx%d)", w, h); 223 224 printf("Using %s\n", aom_codec_iface_name(encoder)); 225 226 // Configuration 227 res = aom_codec_enc_config_default(encoder, &cfg, 0); 228 if (res) die_codec(&codec, "Failed to get default codec config."); 229 230 cfg.g_w = w; 231 cfg.g_h = h; 232 cfg.g_timebase.num = 1; 233 cfg.g_timebase.den = fps; 234 cfg.rc_target_bitrate = bitrate; 235 236 if (!(infile = fopen(infile_arg, "rb"))) 237 die("Failed to open %s for reading", infile_arg); 238 239 // Pass 0 240 cfg.g_pass = AOM_RC_FIRST_PASS; 241 stats = pass0(&raw, infile, encoder, &cfg, limit); 242 243 // Pass 1 244 rewind(infile); 245 cfg.g_pass = AOM_RC_LAST_PASS; 246 cfg.rc_twopass_stats_in = stats; 247 pass1(&raw, infile, outfile_arg, encoder, &cfg, limit); 248 free(stats.buf); 249 250 aom_img_free(&raw); 251 fclose(infile); 252 253 return EXIT_SUCCESS; 254 }