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