]> git.sesse.net Git - ffmpeg/blob - libavcodec/libaomenc.c
Merge commit '3a7b4ae62c798edbd82bcd8fef863c74ed2acd4a'
[ffmpeg] / libavcodec / libaomenc.c
1 /*
2  * Copyright (c) 2010, Google, Inc.
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU Lesser General Public
8  * License as published by the Free Software Foundation; either
9  * version 2.1 of the License, or (at your option) any later version.
10  *
11  * FFmpeg is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14  * Lesser General Public License for more details.
15  *
16  * You should have received a copy of the GNU Lesser General Public
17  * License along with FFmpeg; if not, write to the Free Software
18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19  */
20
21 /**
22  * @file
23  * AV1 encoder support via libaom
24  */
25
26 #define AOM_DISABLE_CTRL_TYPECHECKS 1
27 #include <aom/aom_encoder.h>
28 #include <aom/aomcx.h>
29
30 #include "libavutil/avassert.h"
31 #include "libavutil/base64.h"
32 #include "libavutil/common.h"
33 #include "libavutil/mathematics.h"
34 #include "libavutil/opt.h"
35 #include "libavutil/pixdesc.h"
36
37 #include "avcodec.h"
38 #include "internal.h"
39 #include "profiles.h"
40
41 /*
42  * Portion of struct aom_codec_cx_pkt from aom_encoder.h.
43  * One encoded frame returned from the library.
44  */
45 struct FrameListData {
46     void *buf;                       /**< compressed data buffer */
47     size_t sz;                       /**< length of compressed data */
48     int64_t pts;                     /**< time stamp to show frame
49                                           (in timebase units) */
50     unsigned long duration;          /**< duration to show frame
51                                           (in timebase units) */
52     uint32_t flags;                  /**< flags for this frame */
53     struct FrameListData *next;
54 };
55
56 typedef struct AOMEncoderContext {
57     AVClass *class;
58     struct aom_codec_ctx encoder;
59     struct aom_image rawimg;
60     struct aom_fixed_buf twopass_stats;
61     struct FrameListData *coded_frame_list;
62     int cpu_used;
63     int auto_alt_ref;
64     int lag_in_frames;
65     int error_resilient;
66     int crf;
67     int static_thresh;
68     int drop_threshold;
69     int noise_sensitivity;
70 } AOMContext;
71
72 static const char *const ctlidstr[] = {
73     [AOME_SET_CPUUSED]          = "AOME_SET_CPUUSED",
74     [AOME_SET_CQ_LEVEL]         = "AOME_SET_CQ_LEVEL",
75     [AOME_SET_ENABLEAUTOALTREF] = "AOME_SET_ENABLEAUTOALTREF",
76     [AOME_SET_STATIC_THRESHOLD] = "AOME_SET_STATIC_THRESHOLD",
77     [AV1E_SET_COLOR_RANGE]      = "AV1E_SET_COLOR_RANGE",
78     [AV1E_SET_COLOR_PRIMARIES]  = "AV1E_SET_COLOR_PRIMARIES",
79     [AV1E_SET_MATRIX_COEFFICIENTS] = "AV1E_SET_MATRIX_COEFFICIENTS",
80     [AV1E_SET_TRANSFER_CHARACTERISTICS] = "AV1E_SET_TRANSFER_CHARACTERISTICS",
81 };
82
83 static av_cold void log_encoder_error(AVCodecContext *avctx, const char *desc)
84 {
85     AOMContext *ctx    = avctx->priv_data;
86     const char *error  = aom_codec_error(&ctx->encoder);
87     const char *detail = aom_codec_error_detail(&ctx->encoder);
88
89     av_log(avctx, AV_LOG_ERROR, "%s: %s\n", desc, error);
90     if (detail)
91         av_log(avctx, AV_LOG_ERROR, "  Additional information: %s\n", detail);
92 }
93
94 static av_cold void dump_enc_cfg(AVCodecContext *avctx,
95                                  const struct aom_codec_enc_cfg *cfg)
96 {
97     int width = -30;
98     int level = AV_LOG_DEBUG;
99
100     av_log(avctx, level, "aom_codec_enc_cfg\n");
101     av_log(avctx, level, "generic settings\n"
102                          "  %*s%u\n  %*s%u\n  %*s%u\n  %*s%u\n  %*s%u\n"
103                          "  %*s%u\n  %*s%u\n"
104                          "  %*s{%u/%u}\n  %*s%u\n  %*s%d\n  %*s%u\n",
105            width, "g_usage:",           cfg->g_usage,
106            width, "g_threads:",         cfg->g_threads,
107            width, "g_profile:",         cfg->g_profile,
108            width, "g_w:",               cfg->g_w,
109            width, "g_h:",               cfg->g_h,
110            width, "g_bit_depth:",       cfg->g_bit_depth,
111            width, "g_input_bit_depth:", cfg->g_input_bit_depth,
112            width, "g_timebase:",        cfg->g_timebase.num, cfg->g_timebase.den,
113            width, "g_error_resilient:", cfg->g_error_resilient,
114            width, "g_pass:",            cfg->g_pass,
115            width, "g_lag_in_frames:",   cfg->g_lag_in_frames);
116     av_log(avctx, level, "rate control settings\n"
117                          "  %*s%u\n  %*s%d\n  %*s%p(%"SIZE_SPECIFIER")\n  %*s%u\n",
118            width, "rc_dropframe_thresh:", cfg->rc_dropframe_thresh,
119            width, "rc_end_usage:",        cfg->rc_end_usage,
120            width, "rc_twopass_stats_in:", cfg->rc_twopass_stats_in.buf, cfg->rc_twopass_stats_in.sz,
121            width, "rc_target_bitrate:",   cfg->rc_target_bitrate);
122     av_log(avctx, level, "quantizer settings\n"
123                          "  %*s%u\n  %*s%u\n",
124            width, "rc_min_quantizer:", cfg->rc_min_quantizer,
125            width, "rc_max_quantizer:", cfg->rc_max_quantizer);
126     av_log(avctx, level, "bitrate tolerance\n"
127                          "  %*s%u\n  %*s%u\n",
128            width, "rc_undershoot_pct:", cfg->rc_undershoot_pct,
129            width, "rc_overshoot_pct:",  cfg->rc_overshoot_pct);
130     av_log(avctx, level, "decoder buffer model\n"
131                          "  %*s%u\n  %*s%u\n  %*s%u\n",
132            width, "rc_buf_sz:",         cfg->rc_buf_sz,
133            width, "rc_buf_initial_sz:", cfg->rc_buf_initial_sz,
134            width, "rc_buf_optimal_sz:", cfg->rc_buf_optimal_sz);
135     av_log(avctx, level, "2 pass rate control settings\n"
136                          "  %*s%u\n  %*s%u\n  %*s%u\n",
137            width, "rc_2pass_vbr_bias_pct:",       cfg->rc_2pass_vbr_bias_pct,
138            width, "rc_2pass_vbr_minsection_pct:", cfg->rc_2pass_vbr_minsection_pct,
139            width, "rc_2pass_vbr_maxsection_pct:", cfg->rc_2pass_vbr_maxsection_pct);
140     av_log(avctx, level, "keyframing settings\n"
141                          "  %*s%d\n  %*s%u\n  %*s%u\n",
142            width, "kf_mode:",     cfg->kf_mode,
143            width, "kf_min_dist:", cfg->kf_min_dist,
144            width, "kf_max_dist:", cfg->kf_max_dist);
145     av_log(avctx, level, "\n");
146 }
147
148 static void coded_frame_add(void *list, struct FrameListData *cx_frame)
149 {
150     struct FrameListData **p = list;
151
152     while (*p)
153         p = &(*p)->next;
154     *p = cx_frame;
155     cx_frame->next = NULL;
156 }
157
158 static av_cold void free_coded_frame(struct FrameListData *cx_frame)
159 {
160     av_freep(&cx_frame->buf);
161     av_freep(&cx_frame);
162 }
163
164 static av_cold void free_frame_list(struct FrameListData *list)
165 {
166     struct FrameListData *p = list;
167
168     while (p) {
169         list = list->next;
170         free_coded_frame(p);
171         p = list;
172     }
173 }
174
175 static av_cold int codecctl_int(AVCodecContext *avctx,
176                                 enum aome_enc_control_id id, int val)
177 {
178     AOMContext *ctx = avctx->priv_data;
179     char buf[80];
180     int width = -30;
181     int res;
182
183     snprintf(buf, sizeof(buf), "%s:", ctlidstr[id]);
184     av_log(avctx, AV_LOG_DEBUG, "  %*s%d\n", width, buf, val);
185
186     res = aom_codec_control(&ctx->encoder, id, val);
187     if (res != AOM_CODEC_OK) {
188         snprintf(buf, sizeof(buf), "Failed to set %s codec control",
189                  ctlidstr[id]);
190         log_encoder_error(avctx, buf);
191         return AVERROR(EINVAL);
192     }
193
194     return 0;
195 }
196
197 static av_cold int aom_free(AVCodecContext *avctx)
198 {
199     AOMContext *ctx = avctx->priv_data;
200
201     aom_codec_destroy(&ctx->encoder);
202     av_freep(&ctx->twopass_stats.buf);
203     av_freep(&avctx->stats_out);
204     free_frame_list(ctx->coded_frame_list);
205     return 0;
206 }
207
208 static int set_pix_fmt(AVCodecContext *avctx, aom_codec_caps_t codec_caps,
209                        struct aom_codec_enc_cfg *enccfg, aom_codec_flags_t *flags,
210                        aom_img_fmt_t *img_fmt)
211 {
212     AOMContext av_unused *ctx = avctx->priv_data;
213     enccfg->g_bit_depth = enccfg->g_input_bit_depth = 8;
214     switch (avctx->pix_fmt) {
215     case AV_PIX_FMT_YUV420P:
216         enccfg->g_profile = FF_PROFILE_AV1_MAIN;
217         *img_fmt = AOM_IMG_FMT_I420;
218         return 0;
219     case AV_PIX_FMT_YUV422P:
220         enccfg->g_profile = FF_PROFILE_AV1_PROFESSIONAL;
221         *img_fmt = AOM_IMG_FMT_I422;
222         return 0;
223     case AV_PIX_FMT_GBRP:
224     case AV_PIX_FMT_YUV444P:
225         enccfg->g_profile = FF_PROFILE_AV1_HIGH;
226         *img_fmt = AOM_IMG_FMT_I444;
227         return 0;
228     case AV_PIX_FMT_YUV420P10:
229     case AV_PIX_FMT_YUV420P12:
230         if (codec_caps & AOM_CODEC_CAP_HIGHBITDEPTH) {
231             enccfg->g_bit_depth = enccfg->g_input_bit_depth =
232                 avctx->pix_fmt == AV_PIX_FMT_YUV420P10 ? 10 : 12;
233             enccfg->g_profile =
234                 enccfg->g_bit_depth == 10 ? FF_PROFILE_AV1_MAIN : FF_PROFILE_AV1_PROFESSIONAL;
235             *img_fmt = AOM_IMG_FMT_I42016;
236             *flags |= AOM_CODEC_USE_HIGHBITDEPTH;
237             return 0;
238         }
239         break;
240     case AV_PIX_FMT_YUV422P10:
241     case AV_PIX_FMT_YUV422P12:
242         if (codec_caps & AOM_CODEC_CAP_HIGHBITDEPTH) {
243             enccfg->g_bit_depth = enccfg->g_input_bit_depth =
244                 avctx->pix_fmt == AV_PIX_FMT_YUV422P10 ? 10 : 12;
245             enccfg->g_profile = FF_PROFILE_AV1_PROFESSIONAL;
246             *img_fmt = AOM_IMG_FMT_I42216;
247             *flags |= AOM_CODEC_USE_HIGHBITDEPTH;
248             return 0;
249         }
250         break;
251     case AV_PIX_FMT_GBRP10:
252     case AV_PIX_FMT_GBRP12:
253     case AV_PIX_FMT_YUV444P10:
254     case AV_PIX_FMT_YUV444P12:
255         if (codec_caps & AOM_CODEC_CAP_HIGHBITDEPTH) {
256             enccfg->g_bit_depth = enccfg->g_input_bit_depth =
257                 avctx->pix_fmt == AV_PIX_FMT_YUV444P10 ||
258                 avctx->pix_fmt == AV_PIX_FMT_GBRP10 ? 10 : 12;
259             enccfg->g_profile =
260                 enccfg->g_bit_depth == 10 ? FF_PROFILE_AV1_HIGH : FF_PROFILE_AV1_PROFESSIONAL;
261             *img_fmt = AOM_IMG_FMT_I44416;
262             *flags |= AOM_CODEC_USE_HIGHBITDEPTH;
263             return 0;
264         }
265         break;
266     default:
267         break;
268     }
269     av_log(avctx, AV_LOG_ERROR, "Unsupported pixel format.\n");
270     return AVERROR_INVALIDDATA;
271 }
272
273 static void set_color_range(AVCodecContext *avctx)
274 {
275     enum aom_color_range aom_cr;
276     switch (avctx->color_range) {
277     case AVCOL_RANGE_UNSPECIFIED:
278     case AVCOL_RANGE_MPEG:       aom_cr = AOM_CR_STUDIO_RANGE; break;
279     case AVCOL_RANGE_JPEG:       aom_cr = AOM_CR_FULL_RANGE;   break;
280     default:
281         av_log(avctx, AV_LOG_WARNING, "Unsupported color range (%d)\n",
282                avctx->color_range);
283         return;
284     }
285
286     codecctl_int(avctx, AV1E_SET_COLOR_RANGE, aom_cr);
287 }
288
289 static av_cold int aom_init(AVCodecContext *avctx,
290                             const struct aom_codec_iface *iface)
291 {
292     AOMContext *ctx = avctx->priv_data;
293     struct aom_codec_enc_cfg enccfg = { 0 };
294     aom_codec_flags_t flags = 0;
295     AVCPBProperties *cpb_props;
296     int res;
297     aom_img_fmt_t img_fmt;
298     aom_codec_caps_t codec_caps = aom_codec_get_caps(iface);
299
300     av_log(avctx, AV_LOG_INFO, "%s\n", aom_codec_version_str());
301     av_log(avctx, AV_LOG_VERBOSE, "%s\n", aom_codec_build_config());
302
303     if ((res = aom_codec_enc_config_default(iface, &enccfg, 0)) != AOM_CODEC_OK) {
304         av_log(avctx, AV_LOG_ERROR, "Failed to get config: %s\n",
305                aom_codec_err_to_string(res));
306         return AVERROR(EINVAL);
307     }
308
309     if (set_pix_fmt(avctx, codec_caps, &enccfg, &flags, &img_fmt))
310         return AVERROR(EINVAL);
311
312     if(!avctx->bit_rate)
313         if(avctx->rc_max_rate || avctx->rc_buffer_size || avctx->rc_initial_buffer_occupancy) {
314             av_log( avctx, AV_LOG_ERROR, "Rate control parameters set without a bitrate\n");
315             return AVERROR(EINVAL);
316         }
317
318     dump_enc_cfg(avctx, &enccfg);
319
320     enccfg.g_w            = avctx->width;
321     enccfg.g_h            = avctx->height;
322     enccfg.g_timebase.num = avctx->time_base.num;
323     enccfg.g_timebase.den = avctx->time_base.den;
324     enccfg.g_threads      = avctx->thread_count;
325
326     if (ctx->lag_in_frames >= 0)
327         enccfg.g_lag_in_frames = ctx->lag_in_frames;
328
329     if (avctx->flags & AV_CODEC_FLAG_PASS1)
330         enccfg.g_pass = AOM_RC_FIRST_PASS;
331     else if (avctx->flags & AV_CODEC_FLAG_PASS2)
332         enccfg.g_pass = AOM_RC_LAST_PASS;
333     else
334         enccfg.g_pass = AOM_RC_ONE_PASS;
335
336     if (avctx->rc_min_rate == avctx->rc_max_rate &&
337         avctx->rc_min_rate == avctx->bit_rate && avctx->bit_rate) {
338         enccfg.rc_end_usage = AOM_CBR;
339     } else if (ctx->crf >= 0) {
340         enccfg.rc_end_usage = AOM_CQ;
341         if (!avctx->bit_rate)
342             enccfg.rc_end_usage = AOM_Q;
343     }
344
345     if (avctx->bit_rate) {
346         enccfg.rc_target_bitrate = av_rescale_rnd(avctx->bit_rate, 1, 1000,
347                                                   AV_ROUND_NEAR_INF);
348     } else if (enccfg.rc_end_usage != AOM_Q) {
349         if (enccfg.rc_end_usage == AOM_CQ) {
350             enccfg.rc_target_bitrate = 1000000;
351         } else {
352             avctx->bit_rate = enccfg.rc_target_bitrate * 1000;
353             av_log(avctx, AV_LOG_WARNING,
354                    "Neither bitrate nor constrained quality specified, using default bitrate of %dkbit/sec\n",
355                    enccfg.rc_target_bitrate);
356         }
357     }
358
359     if (avctx->qmin >= 0)
360         enccfg.rc_min_quantizer = avctx->qmin;
361     if (avctx->qmax >= 0)
362         enccfg.rc_max_quantizer = avctx->qmax;
363
364     if (enccfg.rc_end_usage == AOM_CQ || enccfg.rc_end_usage == AOM_Q) {
365         if (ctx->crf < enccfg.rc_min_quantizer || ctx->crf > enccfg.rc_max_quantizer) {
366             av_log(avctx, AV_LOG_ERROR,
367                    "CQ level %d must be between minimum and maximum quantizer value (%d-%d)\n",
368                    ctx->crf, enccfg.rc_min_quantizer, enccfg.rc_max_quantizer);
369             return AVERROR(EINVAL);
370         }
371     }
372
373     enccfg.rc_dropframe_thresh = ctx->drop_threshold;
374
375     // 0-100 (0 => CBR, 100 => VBR)
376     enccfg.rc_2pass_vbr_bias_pct       = round(avctx->qcompress * 100);
377     if (avctx->bit_rate)
378         enccfg.rc_2pass_vbr_minsection_pct =
379             avctx->rc_min_rate * 100LL / avctx->bit_rate;
380     if (avctx->rc_max_rate)
381         enccfg.rc_2pass_vbr_maxsection_pct =
382             avctx->rc_max_rate * 100LL / avctx->bit_rate;
383
384     if (avctx->rc_buffer_size)
385         enccfg.rc_buf_sz =
386             avctx->rc_buffer_size * 1000LL / avctx->bit_rate;
387     if (avctx->rc_initial_buffer_occupancy)
388         enccfg.rc_buf_initial_sz =
389             avctx->rc_initial_buffer_occupancy * 1000LL / avctx->bit_rate;
390     enccfg.rc_buf_optimal_sz = enccfg.rc_buf_sz * 5 / 6;
391
392     // _enc_init() will balk if kf_min_dist differs from max w/AOM_KF_AUTO
393     if (avctx->keyint_min >= 0 && avctx->keyint_min == avctx->gop_size)
394         enccfg.kf_min_dist = avctx->keyint_min;
395     if (avctx->gop_size >= 0)
396         enccfg.kf_max_dist = avctx->gop_size;
397
398     if (enccfg.g_pass == AOM_RC_FIRST_PASS)
399         enccfg.g_lag_in_frames = 0;
400     else if (enccfg.g_pass == AOM_RC_LAST_PASS) {
401         int decode_size, ret;
402
403         if (!avctx->stats_in) {
404             av_log(avctx, AV_LOG_ERROR, "No stats file for second pass\n");
405             return AVERROR_INVALIDDATA;
406         }
407
408         ctx->twopass_stats.sz = strlen(avctx->stats_in) * 3 / 4;
409         ret                   = av_reallocp(&ctx->twopass_stats.buf, ctx->twopass_stats.sz);
410         if (ret < 0) {
411             av_log(avctx, AV_LOG_ERROR,
412                    "Stat buffer alloc (%"SIZE_SPECIFIER" bytes) failed\n",
413                    ctx->twopass_stats.sz);
414             ctx->twopass_stats.sz = 0;
415             return ret;
416         }
417         decode_size = av_base64_decode(ctx->twopass_stats.buf, avctx->stats_in,
418                                        ctx->twopass_stats.sz);
419         if (decode_size < 0) {
420             av_log(avctx, AV_LOG_ERROR, "Stat buffer decode failed\n");
421             return AVERROR_INVALIDDATA;
422         }
423
424         ctx->twopass_stats.sz      = decode_size;
425         enccfg.rc_twopass_stats_in = ctx->twopass_stats;
426     }
427
428     /* 0-3: For non-zero values the encoder increasingly optimizes for reduced
429      * complexity playback on low powered devices at the expense of encode
430      * quality. */
431     if (avctx->profile != FF_PROFILE_UNKNOWN)
432         enccfg.g_profile = avctx->profile;
433
434     enccfg.g_error_resilient = ctx->error_resilient;
435
436     dump_enc_cfg(avctx, &enccfg);
437     /* Construct Encoder Context */
438     res = aom_codec_enc_init(&ctx->encoder, iface, &enccfg, flags);
439     if (res != AOM_CODEC_OK) {
440         log_encoder_error(avctx, "Failed to initialize encoder");
441         return AVERROR(EINVAL);
442     }
443
444     // codec control failures are currently treated only as warnings
445     av_log(avctx, AV_LOG_DEBUG, "aom_codec_control\n");
446     codecctl_int(avctx, AOME_SET_CPUUSED, ctx->cpu_used);
447     if (ctx->auto_alt_ref >= 0)
448         codecctl_int(avctx, AOME_SET_ENABLEAUTOALTREF, ctx->auto_alt_ref);
449
450     codecctl_int(avctx, AOME_SET_STATIC_THRESHOLD, ctx->static_thresh);
451     if (ctx->crf >= 0)
452         codecctl_int(avctx, AOME_SET_CQ_LEVEL,          ctx->crf);
453
454     codecctl_int(avctx, AV1E_SET_COLOR_PRIMARIES, avctx->color_primaries);
455     codecctl_int(avctx, AV1E_SET_MATRIX_COEFFICIENTS, avctx->colorspace);
456     codecctl_int(avctx, AV1E_SET_TRANSFER_CHARACTERISTICS, avctx->color_trc);
457     set_color_range(avctx);
458
459     // provide dummy value to initialize wrapper, values will be updated each _encode()
460     aom_img_wrap(&ctx->rawimg, img_fmt, avctx->width, avctx->height, 1,
461                  (unsigned char*)1);
462
463     if (codec_caps & AOM_CODEC_CAP_HIGHBITDEPTH)
464         ctx->rawimg.bit_depth = enccfg.g_bit_depth;
465
466     cpb_props = ff_add_cpb_side_data(avctx);
467     if (!cpb_props)
468         return AVERROR(ENOMEM);
469
470     if (enccfg.rc_end_usage == AOM_CBR ||
471         enccfg.g_pass != AOM_RC_ONE_PASS) {
472         cpb_props->max_bitrate = avctx->rc_max_rate;
473         cpb_props->min_bitrate = avctx->rc_min_rate;
474         cpb_props->avg_bitrate = avctx->bit_rate;
475     }
476     cpb_props->buffer_size = avctx->rc_buffer_size;
477
478     return 0;
479 }
480
481 static inline void cx_pktcpy(struct FrameListData *dst,
482                              const struct aom_codec_cx_pkt *src)
483 {
484     dst->pts      = src->data.frame.pts;
485     dst->duration = src->data.frame.duration;
486     dst->flags    = src->data.frame.flags;
487     dst->sz       = src->data.frame.sz;
488     dst->buf      = src->data.frame.buf;
489 }
490
491 /**
492  * Store coded frame information in format suitable for return from encode2().
493  *
494  * Write information from @a cx_frame to @a pkt
495  * @return packet data size on success
496  * @return a negative AVERROR on error
497  */
498 static int storeframe(AVCodecContext *avctx, struct FrameListData *cx_frame,
499                       AVPacket *pkt)
500 {
501     int ret = ff_alloc_packet2(avctx, pkt, cx_frame->sz, 0);
502     if (ret < 0) {
503         av_log(avctx, AV_LOG_ERROR,
504                "Error getting output packet of size %"SIZE_SPECIFIER".\n", cx_frame->sz);
505         return ret;
506     }
507     memcpy(pkt->data, cx_frame->buf, pkt->size);
508     pkt->pts = pkt->dts = cx_frame->pts;
509
510     if (!!(cx_frame->flags & AOM_FRAME_IS_KEY))
511         pkt->flags |= AV_PKT_FLAG_KEY;
512     return pkt->size;
513 }
514
515 /**
516  * Queue multiple output frames from the encoder, returning the front-most.
517  * In cases where aom_codec_get_cx_data() returns more than 1 frame append
518  * the frame queue. Return the head frame if available.
519  * @return Stored frame size
520  * @return AVERROR(EINVAL) on output size error
521  * @return AVERROR(ENOMEM) on coded frame queue data allocation error
522  */
523 static int queue_frames(AVCodecContext *avctx, AVPacket *pkt_out)
524 {
525     AOMContext *ctx = avctx->priv_data;
526     const struct aom_codec_cx_pkt *pkt;
527     const void *iter = NULL;
528     int size = 0;
529
530     if (ctx->coded_frame_list) {
531         struct FrameListData *cx_frame = ctx->coded_frame_list;
532         /* return the leading frame if we've already begun queueing */
533         size = storeframe(avctx, cx_frame, pkt_out);
534         if (size < 0)
535             return size;
536         ctx->coded_frame_list = cx_frame->next;
537         free_coded_frame(cx_frame);
538     }
539
540     /* consume all available output from the encoder before returning. buffers
541      * are only good through the next aom_codec call */
542     while ((pkt = aom_codec_get_cx_data(&ctx->encoder, &iter))) {
543         switch (pkt->kind) {
544         case AOM_CODEC_CX_FRAME_PKT:
545             if (!size) {
546                 struct FrameListData cx_frame;
547
548                 /* avoid storing the frame when the list is empty and we haven't yet
549                  * provided a frame for output */
550                 av_assert0(!ctx->coded_frame_list);
551                 cx_pktcpy(&cx_frame, pkt);
552                 size = storeframe(avctx, &cx_frame, pkt_out);
553                 if (size < 0)
554                     return size;
555             } else {
556                 struct FrameListData *cx_frame =
557                     av_malloc(sizeof(struct FrameListData));
558
559                 if (!cx_frame) {
560                     av_log(avctx, AV_LOG_ERROR,
561                            "Frame queue element alloc failed\n");
562                     return AVERROR(ENOMEM);
563                 }
564                 cx_pktcpy(cx_frame, pkt);
565                 cx_frame->buf = av_malloc(cx_frame->sz);
566
567                 if (!cx_frame->buf) {
568                     av_log(avctx, AV_LOG_ERROR,
569                            "Data buffer alloc (%"SIZE_SPECIFIER" bytes) failed\n",
570                            cx_frame->sz);
571                     av_freep(&cx_frame);
572                     return AVERROR(ENOMEM);
573                 }
574                 memcpy(cx_frame->buf, pkt->data.frame.buf, pkt->data.frame.sz);
575                 coded_frame_add(&ctx->coded_frame_list, cx_frame);
576             }
577             break;
578         case AOM_CODEC_STATS_PKT:
579         {
580             struct aom_fixed_buf *stats = &ctx->twopass_stats;
581             int err;
582             if ((err = av_reallocp(&stats->buf,
583                                    stats->sz +
584                                    pkt->data.twopass_stats.sz)) < 0) {
585                 stats->sz = 0;
586                 av_log(avctx, AV_LOG_ERROR, "Stat buffer realloc failed\n");
587                 return err;
588             }
589             memcpy((uint8_t *)stats->buf + stats->sz,
590                    pkt->data.twopass_stats.buf, pkt->data.twopass_stats.sz);
591             stats->sz += pkt->data.twopass_stats.sz;
592             break;
593         }
594         case AOM_CODEC_PSNR_PKT: // FIXME add support for AV_CODEC_FLAG_PSNR
595         case AOM_CODEC_CUSTOM_PKT:
596             // ignore unsupported/unrecognized packet types
597             break;
598         }
599     }
600
601     return size;
602 }
603
604 static int aom_encode(AVCodecContext *avctx, AVPacket *pkt,
605                       const AVFrame *frame, int *got_packet)
606 {
607     AOMContext *ctx = avctx->priv_data;
608     struct aom_image *rawimg = NULL;
609     int64_t timestamp = 0;
610     int res, coded_size;
611     aom_enc_frame_flags_t flags = 0;
612
613     if (frame) {
614         rawimg                      = &ctx->rawimg;
615         rawimg->planes[AOM_PLANE_Y] = frame->data[0];
616         rawimg->planes[AOM_PLANE_U] = frame->data[1];
617         rawimg->planes[AOM_PLANE_V] = frame->data[2];
618         rawimg->stride[AOM_PLANE_Y] = frame->linesize[0];
619         rawimg->stride[AOM_PLANE_U] = frame->linesize[1];
620         rawimg->stride[AOM_PLANE_V] = frame->linesize[2];
621         timestamp                   = frame->pts;
622         switch (frame->color_range) {
623         case AVCOL_RANGE_MPEG:
624             rawimg->range = AOM_CR_STUDIO_RANGE;
625             break;
626         case AVCOL_RANGE_JPEG:
627             rawimg->range = AOM_CR_FULL_RANGE;
628             break;
629         }
630
631         if (frame->pict_type == AV_PICTURE_TYPE_I)
632             flags |= AOM_EFLAG_FORCE_KF;
633     }
634
635     res = aom_codec_encode(&ctx->encoder, rawimg, timestamp,
636                            avctx->ticks_per_frame, flags);
637     if (res != AOM_CODEC_OK) {
638         log_encoder_error(avctx, "Error encoding frame");
639         return AVERROR_INVALIDDATA;
640     }
641     coded_size = queue_frames(avctx, pkt);
642
643     if (!frame && avctx->flags & AV_CODEC_FLAG_PASS1) {
644         size_t b64_size = AV_BASE64_SIZE(ctx->twopass_stats.sz);
645
646         avctx->stats_out = av_malloc(b64_size);
647         if (!avctx->stats_out) {
648             av_log(avctx, AV_LOG_ERROR, "Stat buffer alloc (%"SIZE_SPECIFIER" bytes) failed\n",
649                    b64_size);
650             return AVERROR(ENOMEM);
651         }
652         av_base64_encode(avctx->stats_out, b64_size, ctx->twopass_stats.buf,
653                          ctx->twopass_stats.sz);
654     }
655
656     *got_packet = !!coded_size;
657     return 0;
658 }
659
660 static const enum AVPixelFormat av1_pix_fmts[] = {
661     AV_PIX_FMT_YUV420P,
662     AV_PIX_FMT_YUV422P,
663     AV_PIX_FMT_YUV444P,
664     AV_PIX_FMT_GBRP,
665     AV_PIX_FMT_NONE
666 };
667
668 static const enum AVPixelFormat av1_pix_fmts_highbd[] = {
669     AV_PIX_FMT_YUV420P,
670     AV_PIX_FMT_YUV422P,
671     AV_PIX_FMT_YUV444P,
672     AV_PIX_FMT_YUV420P10,
673     AV_PIX_FMT_YUV422P10,
674     AV_PIX_FMT_YUV444P10,
675     AV_PIX_FMT_YUV420P12,
676     AV_PIX_FMT_YUV422P12,
677     AV_PIX_FMT_YUV444P12,
678     AV_PIX_FMT_GBRP,
679     AV_PIX_FMT_GBRP10,
680     AV_PIX_FMT_GBRP12,
681     AV_PIX_FMT_NONE
682 };
683
684 static av_cold void av1_init_static(AVCodec *codec)
685 {
686     aom_codec_caps_t codec_caps = aom_codec_get_caps(aom_codec_av1_cx());
687     if (codec_caps & AOM_CODEC_CAP_HIGHBITDEPTH)
688         codec->pix_fmts = av1_pix_fmts_highbd;
689     else
690         codec->pix_fmts = av1_pix_fmts;
691 }
692
693 static av_cold int av1_init(AVCodecContext *avctx)
694 {
695     return aom_init(avctx, aom_codec_av1_cx());
696 }
697
698 #define OFFSET(x) offsetof(AOMContext, x)
699 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
700 static const AVOption options[] = {
701     { "cpu-used",        "Quality/Speed ratio modifier",           OFFSET(cpu_used),        AV_OPT_TYPE_INT, {.i64 = 1}, -8, 8, VE},
702     { "auto-alt-ref",    "Enable use of alternate reference "
703                          "frames (2-pass only)",                   OFFSET(auto_alt_ref),    AV_OPT_TYPE_INT, {.i64 = -1},      -1,      2,       VE},
704     { "lag-in-frames",   "Number of frames to look ahead at for "
705                          "alternate reference frame selection",    OFFSET(lag_in_frames),   AV_OPT_TYPE_INT, {.i64 = -1},      -1,      INT_MAX, VE},
706     { "error-resilience", "Error resilience configuration", OFFSET(error_resilient), AV_OPT_TYPE_FLAGS, {.i64 = 0}, INT_MIN, INT_MAX, VE, "er"},
707     { "default",         "Improve resiliency against losses of whole frames", 0, AV_OPT_TYPE_CONST, {.i64 = AOM_ERROR_RESILIENT_DEFAULT}, 0, 0, VE, "er"},
708     { "partitions",      "The frame partitions are independently decodable "
709                          "by the bool decoder, meaning that partitions can be decoded even "
710                          "though earlier partitions have been lost. Note that intra predicition"
711                          " is still done over the partition boundary.",       0, AV_OPT_TYPE_CONST, {.i64 = AOM_ERROR_RESILIENT_PARTITIONS}, 0, 0, VE, "er"},
712     { "crf",              "Select the quality for constant quality mode", offsetof(AOMContext, crf), AV_OPT_TYPE_INT, {.i64 = -1}, -1, 63, VE },
713     { "static-thresh",    "A change threshold on blocks below which they will be skipped by the encoder", OFFSET(static_thresh), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INT_MAX, VE },
714     { "drop-threshold",   "Frame drop threshold", offsetof(AOMContext, drop_threshold), AV_OPT_TYPE_INT, {.i64 = 0 }, INT_MIN, INT_MAX, VE },
715     { "noise-sensitivity", "Noise sensitivity", OFFSET(noise_sensitivity), AV_OPT_TYPE_INT, {.i64 = 0 }, 0, 4, VE},
716     { NULL }
717 };
718
719 static const AVCodecDefault defaults[] = {
720     { "qmin",             "-1" },
721     { "qmax",             "-1" },
722     { "g",                "-1" },
723     { "keyint_min",       "-1" },
724     { NULL },
725 };
726
727 static const AVClass class_aom = {
728     .class_name = "libaom-av1 encoder",
729     .item_name  = av_default_item_name,
730     .option     = options,
731     .version    = LIBAVUTIL_VERSION_INT,
732 };
733
734 AVCodec ff_libaom_av1_encoder = {
735     .name           = "libaom-av1",
736     .long_name      = NULL_IF_CONFIG_SMALL("libaom AV1"),
737     .type           = AVMEDIA_TYPE_VIDEO,
738     .id             = AV_CODEC_ID_AV1,
739     .priv_data_size = sizeof(AOMContext),
740     .init           = av1_init,
741     .encode2        = aom_encode,
742     .close          = aom_free,
743     .capabilities   = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AUTO_THREADS | AV_CODEC_CAP_EXPERIMENTAL,
744     .profiles       = NULL_IF_CONFIG_SMALL(ff_av1_profiles),
745     .priv_class     = &class_aom,
746     .defaults       = defaults,
747     .init_static_data = av1_init_static,
748     .wrapper_name   = "libaom",
749 };