4 * Copyright (c) 2019 Derek Buitenhuis
6 * This file is part of FFmpeg.
8 * FFmpeg is free software; you can redistribute it and/or
9 * modify it under the terms of the GNU Lesser General Public
10 * License as published by the Free Software Foundation; either
11 * version 2.1 of the License, or (at your option) any later version.
13 * FFmpeg is distributed in the hope that it will be useful,
14 * but WITHOUT ANY WARRANTY; without even the implied warranty of
15 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
16 * Lesser General Public License for more details.
18 * You should have received a copy of the GNU Lesser General Public
19 * License along with FFmpeg; if not, write to the Free Software
20 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
25 #include "libavutil/internal.h"
26 #include "libavutil/avassert.h"
27 #include "libavutil/base64.h"
28 #include "libavutil/common.h"
29 #include "libavutil/mathematics.h"
30 #include "libavutil/opt.h"
31 #include "libavutil/pixdesc.h"
36 typedef struct librav1eContext {
48 AVDictionary *rav1e_opts;
56 static inline RaPixelRange range_map(enum AVPixelFormat pix_fmt, enum AVColorRange range)
59 case AV_PIX_FMT_YUVJ420P:
60 case AV_PIX_FMT_YUVJ422P:
61 case AV_PIX_FMT_YUVJ444P:
62 return RA_PIXEL_RANGE_FULL;
66 case AVCOL_RANGE_JPEG:
67 return RA_PIXEL_RANGE_FULL;
68 case AVCOL_RANGE_MPEG:
70 return RA_PIXEL_RANGE_LIMITED;
74 static inline RaChromaSampling pix_fmt_map(enum AVPixelFormat pix_fmt)
77 case AV_PIX_FMT_YUV420P:
78 case AV_PIX_FMT_YUVJ420P:
79 case AV_PIX_FMT_YUV420P10:
80 case AV_PIX_FMT_YUV420P12:
81 return RA_CHROMA_SAMPLING_CS420;
82 case AV_PIX_FMT_YUV422P:
83 case AV_PIX_FMT_YUVJ422P:
84 case AV_PIX_FMT_YUV422P10:
85 case AV_PIX_FMT_YUV422P12:
86 return RA_CHROMA_SAMPLING_CS422;
87 case AV_PIX_FMT_YUV444P:
88 case AV_PIX_FMT_YUVJ444P:
89 case AV_PIX_FMT_YUV444P10:
90 case AV_PIX_FMT_YUV444P12:
91 return RA_CHROMA_SAMPLING_CS444;
97 static inline RaChromaSamplePosition chroma_loc_map(enum AVChromaLocation chroma_loc)
100 case AVCHROMA_LOC_LEFT:
101 return RA_CHROMA_SAMPLE_POSITION_VERTICAL;
102 case AVCHROMA_LOC_TOPLEFT:
103 return RA_CHROMA_SAMPLE_POSITION_COLOCATED;
105 return RA_CHROMA_SAMPLE_POSITION_UNKNOWN;
109 static int get_stats(AVCodecContext *avctx, int eos)
111 librav1eContext *ctx = avctx->priv_data;
112 RaData* buf = rav1e_twopass_out(ctx->ctx);
117 uint8_t *tmp = av_fast_realloc(ctx->pass_data, &ctx->pass_size,
118 ctx->pass_pos + buf->len);
120 rav1e_data_unref(buf);
121 return AVERROR(ENOMEM);
124 ctx->pass_data = tmp;
125 memcpy(ctx->pass_data + ctx->pass_pos, buf->data, buf->len);
126 ctx->pass_pos += buf->len;
128 size_t b64_size = AV_BASE64_SIZE(ctx->pass_pos);
130 memcpy(ctx->pass_data, buf->data, buf->len);
132 avctx->stats_out = av_malloc(b64_size);
133 if (!avctx->stats_out) {
134 rav1e_data_unref(buf);
135 return AVERROR(ENOMEM);
138 av_base64_encode(avctx->stats_out, b64_size, ctx->pass_data, ctx->pass_pos);
140 av_freep(&ctx->pass_data);
143 rav1e_data_unref(buf);
148 static int set_stats(AVCodecContext *avctx)
150 librav1eContext *ctx = avctx->priv_data;
153 while (ret > 0 && ctx->pass_size - ctx->pass_pos > 0) {
154 ret = rav1e_twopass_in(ctx->ctx, ctx->pass_data + ctx->pass_pos, ctx->pass_size);
156 return AVERROR_EXTERNAL;
157 ctx->pass_pos += ret;
163 static av_cold int librav1e_encode_close(AVCodecContext *avctx)
165 librav1eContext *ctx = avctx->priv_data;
168 rav1e_context_unref(ctx->ctx);
172 rav1e_frame_unref(ctx->rframe);
176 av_frame_free(&ctx->frame);
177 av_bsf_free(&ctx->bsf);
178 av_freep(&ctx->pass_data);
183 static av_cold int librav1e_encode_init(AVCodecContext *avctx)
185 librav1eContext *ctx = avctx->priv_data;
186 const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(avctx->pix_fmt);
187 RaConfig *cfg = NULL;
191 ctx->frame = av_frame_alloc();
193 return AVERROR(ENOMEM);
195 cfg = rav1e_config_default();
197 av_log(avctx, AV_LOG_ERROR, "Could not allocate rav1e config.\n");
198 return AVERROR_EXTERNAL;
202 * Rav1e currently uses the time base given to it only for ratecontrol... where
203 * the inverse is taken and used as a framerate. So, do what we do in other wrappers
204 * and use the framerate if we can.
206 if (avctx->framerate.num > 0 && avctx->framerate.den > 0) {
207 rav1e_config_set_time_base(cfg, (RaRational) {
208 avctx->framerate.den, avctx->framerate.num
211 rav1e_config_set_time_base(cfg, (RaRational) {
212 avctx->time_base.num * avctx->ticks_per_frame,
217 if ((avctx->flags & AV_CODEC_FLAG_PASS1 || avctx->flags & AV_CODEC_FLAG_PASS2) && !avctx->bit_rate) {
218 av_log(avctx, AV_LOG_ERROR, "A bitrate must be set to use two pass mode.\n");
219 ret = AVERROR_INVALIDDATA;
223 if (avctx->flags & AV_CODEC_FLAG_PASS2) {
224 if (!avctx->stats_in) {
225 av_log(avctx, AV_LOG_ERROR, "No stats file provided for second pass.\n");
226 ret = AVERROR(EINVAL);
230 ctx->pass_size = (strlen(avctx->stats_in) * 3) / 4;
231 ctx->pass_data = av_malloc(ctx->pass_size);
232 if (!ctx->pass_data) {
233 av_log(avctx, AV_LOG_ERROR, "Could not allocate stats buffer.\n");
234 ret = AVERROR(ENOMEM);
238 ctx->pass_size = av_base64_decode(ctx->pass_data, avctx->stats_in, ctx->pass_size);
239 if (ctx->pass_size < 0) {
240 av_log(avctx, AV_LOG_ERROR, "Invalid pass file.\n");
241 ret = AVERROR(EINVAL);
246 if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
247 const AVBitStreamFilter *filter = av_bsf_get_by_name("extract_extradata");
251 av_log(avctx, AV_LOG_ERROR, "extract_extradata bitstream filter "
252 "not found. This is a bug, please report it.\n");
257 bret = av_bsf_alloc(filter, &ctx->bsf);
263 bret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx);
269 bret = av_bsf_init(ctx->bsf);
277 AVDictionaryEntry *en = NULL;
278 while ((en = av_dict_get(ctx->rav1e_opts, "", en, AV_DICT_IGNORE_SUFFIX))) {
279 int parse_ret = rav1e_config_parse(cfg, en->key, en->value);
281 av_log(avctx, AV_LOG_WARNING, "Invalid value for %s: %s.\n", en->key, en->value);
285 rret = rav1e_config_parse_int(cfg, "width", avctx->width);
287 av_log(avctx, AV_LOG_ERROR, "Invalid width passed to rav1e.\n");
288 ret = AVERROR_INVALIDDATA;
292 rret = rav1e_config_parse_int(cfg, "height", avctx->height);
294 av_log(avctx, AV_LOG_ERROR, "Invalid height passed to rav1e.\n");
295 ret = AVERROR_INVALIDDATA;
299 rret = rav1e_config_parse_int(cfg, "threads", avctx->thread_count);
301 av_log(avctx, AV_LOG_WARNING, "Invalid number of threads, defaulting to auto.\n");
303 if (ctx->speed >= 0) {
304 rret = rav1e_config_parse_int(cfg, "speed", ctx->speed);
306 av_log(avctx, AV_LOG_ERROR, "Could not set speed preset.\n");
307 ret = AVERROR_EXTERNAL;
312 /* rav1e handles precedence between 'tiles' and cols/rows for us. */
313 if (ctx->tiles > 0) {
314 rret = rav1e_config_parse_int(cfg, "tiles", ctx->tiles);
316 av_log(avctx, AV_LOG_ERROR, "Could not set number of tiles to encode with.\n");
317 ret = AVERROR_EXTERNAL;
321 if (ctx->tile_rows > 0) {
322 rret = rav1e_config_parse_int(cfg, "tile_rows", ctx->tile_rows);
324 av_log(avctx, AV_LOG_ERROR, "Could not set number of tile rows to encode with.\n");
325 ret = AVERROR_EXTERNAL;
329 if (ctx->tile_cols > 0) {
330 rret = rav1e_config_parse_int(cfg, "tile_cols", ctx->tile_cols);
332 av_log(avctx, AV_LOG_ERROR, "Could not set number of tile cols to encode with.\n");
333 ret = AVERROR_EXTERNAL;
338 if (avctx->gop_size > 0) {
339 rret = rav1e_config_parse_int(cfg, "key_frame_interval", avctx->gop_size);
341 av_log(avctx, AV_LOG_ERROR, "Could not set max keyint.\n");
342 ret = AVERROR_EXTERNAL;
347 if (avctx->keyint_min > 0) {
348 rret = rav1e_config_parse_int(cfg, "min_key_frame_interval", avctx->keyint_min);
350 av_log(avctx, AV_LOG_ERROR, "Could not set min keyint.\n");
351 ret = AVERROR_EXTERNAL;
356 if (avctx->bit_rate && ctx->quantizer < 0) {
357 int max_quantizer = avctx->qmax >= 0 ? avctx->qmax : 255;
359 rret = rav1e_config_parse_int(cfg, "quantizer", max_quantizer);
361 av_log(avctx, AV_LOG_ERROR, "Could not set max quantizer.\n");
362 ret = AVERROR_EXTERNAL;
366 if (avctx->qmin >= 0) {
367 rret = rav1e_config_parse_int(cfg, "min_quantizer", avctx->qmin);
369 av_log(avctx, AV_LOG_ERROR, "Could not set min quantizer.\n");
370 ret = AVERROR_EXTERNAL;
375 rret = rav1e_config_parse_int(cfg, "bitrate", avctx->bit_rate);
377 av_log(avctx, AV_LOG_ERROR, "Could not set bitrate.\n");
378 ret = AVERROR_INVALIDDATA;
381 } else if (ctx->quantizer >= 0) {
383 av_log(avctx, AV_LOG_WARNING, "Both bitrate and quantizer specified. Using quantizer mode.");
385 rret = rav1e_config_parse_int(cfg, "quantizer", ctx->quantizer);
387 av_log(avctx, AV_LOG_ERROR, "Could not set quantizer.\n");
388 ret = AVERROR_EXTERNAL;
393 rret = rav1e_config_set_pixel_format(cfg, desc->comp[0].depth,
394 pix_fmt_map(avctx->pix_fmt),
395 chroma_loc_map(avctx->chroma_sample_location),
396 range_map(avctx->pix_fmt, avctx->color_range));
398 av_log(avctx, AV_LOG_ERROR, "Failed to set pixel format properties.\n");
399 ret = AVERROR_INVALIDDATA;
403 /* rav1e's colorspace enums match standard values. */
404 rret = rav1e_config_set_color_description(cfg, (RaMatrixCoefficients) avctx->colorspace,
405 (RaColorPrimaries) avctx->color_primaries,
406 (RaTransferCharacteristics) avctx->color_trc);
408 av_log(avctx, AV_LOG_WARNING, "Failed to set color properties.\n");
409 if (avctx->err_recognition & AV_EF_EXPLODE) {
410 ret = AVERROR_INVALIDDATA;
415 ctx->ctx = rav1e_context_new(cfg);
417 av_log(avctx, AV_LOG_ERROR, "Failed to create rav1e encode context.\n");
418 ret = AVERROR_EXTERNAL;
426 rav1e_config_unref(cfg);
431 static int librav1e_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
433 librav1eContext *ctx = avctx->priv_data;
434 RaFrame *rframe = ctx->rframe;
435 RaPacket *rpkt = NULL;
439 AVFrame *frame = ctx->frame;
441 ret = ff_encode_get_frame(avctx, frame);
442 if (ret < 0 && ret != AVERROR_EOF)
446 const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(frame->format);
448 rframe = rav1e_frame_new(ctx->ctx);
450 av_log(avctx, AV_LOG_ERROR, "Could not allocate new rav1e frame.\n");
451 av_frame_unref(frame);
452 return AVERROR(ENOMEM);
455 for (int i = 0; i < desc->nb_components; i++) {
456 int shift = i ? desc->log2_chroma_h : 0;
457 int bytes = desc->comp[0].depth == 8 ? 1 : 2;
458 rav1e_frame_fill_plane(rframe, i, frame->data[i],
459 (frame->height >> shift) * frame->linesize[i],
460 frame->linesize[i], bytes);
462 av_frame_unref(frame);
466 ret = rav1e_send_frame(ctx->ctx, rframe);
468 if (ret == RA_ENCODER_STATUS_ENOUGH_DATA) {
469 ctx->rframe = rframe; /* Queue is full. Store the RaFrame to retry next call */
471 rav1e_frame_unref(rframe); /* No need to unref if flushing. */
476 case RA_ENCODER_STATUS_SUCCESS:
477 case RA_ENCODER_STATUS_ENOUGH_DATA:
479 case RA_ENCODER_STATUS_FAILURE:
480 av_log(avctx, AV_LOG_ERROR, "Could not send frame: %s\n", rav1e_status_to_str(ret));
481 return AVERROR_EXTERNAL;
483 av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_send_frame: %s\n", ret, rav1e_status_to_str(ret));
484 return AVERROR_UNKNOWN;
489 if (avctx->flags & AV_CODEC_FLAG_PASS1) {
490 int sret = get_stats(avctx, 0);
493 } else if (avctx->flags & AV_CODEC_FLAG_PASS2) {
494 int sret = set_stats(avctx);
499 ret = rav1e_receive_packet(ctx->ctx, &rpkt);
501 case RA_ENCODER_STATUS_SUCCESS:
503 case RA_ENCODER_STATUS_LIMIT_REACHED:
504 if (avctx->flags & AV_CODEC_FLAG_PASS1) {
505 int sret = get_stats(avctx, 1);
510 case RA_ENCODER_STATUS_ENCODED:
512 case RA_ENCODER_STATUS_NEED_MORE_DATA:
513 if (avctx->internal->draining) {
514 av_log(avctx, AV_LOG_ERROR, "Unexpected error when receiving packet after EOF.\n");
515 return AVERROR_EXTERNAL;
517 return AVERROR(EAGAIN);
518 case RA_ENCODER_STATUS_FAILURE:
519 av_log(avctx, AV_LOG_ERROR, "Could not encode frame: %s\n", rav1e_status_to_str(ret));
520 return AVERROR_EXTERNAL;
522 av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_receive_packet: %s\n", ret, rav1e_status_to_str(ret));
523 return AVERROR_UNKNOWN;
526 ret = av_new_packet(pkt, rpkt->len);
528 av_log(avctx, AV_LOG_ERROR, "Could not allocate packet.\n");
529 rav1e_packet_unref(rpkt);
533 memcpy(pkt->data, rpkt->data, rpkt->len);
535 if (rpkt->frame_type == RA_FRAME_TYPE_KEY)
536 pkt->flags |= AV_PKT_FLAG_KEY;
538 pkt->pts = pkt->dts = rpkt->input_frameno * avctx->ticks_per_frame;
539 rav1e_packet_unref(rpkt);
541 if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
542 int ret = av_bsf_send_packet(ctx->bsf, pkt);
544 av_log(avctx, AV_LOG_ERROR, "extradata extraction send failed.\n");
545 av_packet_unref(pkt);
549 ret = av_bsf_receive_packet(ctx->bsf, pkt);
551 av_log(avctx, AV_LOG_ERROR, "extradata extraction receive failed.\n");
552 av_packet_unref(pkt);
560 #define OFFSET(x) offsetof(librav1eContext, x)
561 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
563 static const AVOption options[] = {
564 { "qp", "use constant quantizer mode", OFFSET(quantizer), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 255, VE },
565 { "speed", "what speed preset to use", OFFSET(speed), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 10, VE },
566 { "tiles", "number of tiles encode with", OFFSET(tiles), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
567 { "tile-rows", "number of tiles rows to encode with", OFFSET(tile_rows), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
568 { "tile-columns", "number of tiles columns to encode with", OFFSET(tile_cols), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
569 { "rav1e-params", "set the rav1e configuration using a :-separated list of key=value parameters", OFFSET(rav1e_opts), AV_OPT_TYPE_DICT, { 0 }, 0, 0, VE },
573 static const AVCodecDefault librav1e_defaults[] = {
576 { "keyint_min", "0" },
582 const enum AVPixelFormat librav1e_pix_fmts[] = {
585 AV_PIX_FMT_YUV420P10,
586 AV_PIX_FMT_YUV420P12,
589 AV_PIX_FMT_YUV422P10,
590 AV_PIX_FMT_YUV422P12,
593 AV_PIX_FMT_YUV444P10,
594 AV_PIX_FMT_YUV444P12,
598 static const AVClass class = {
599 .class_name = "librav1e",
600 .item_name = av_default_item_name,
602 .version = LIBAVUTIL_VERSION_INT,
605 AVCodec ff_librav1e_encoder = {
607 .long_name = NULL_IF_CONFIG_SMALL("librav1e AV1"),
608 .type = AVMEDIA_TYPE_VIDEO,
609 .id = AV_CODEC_ID_AV1,
610 .init = librav1e_encode_init,
611 .receive_packet = librav1e_receive_packet,
612 .close = librav1e_encode_close,
613 .priv_data_size = sizeof(librav1eContext),
614 .priv_class = &class,
615 .defaults = librav1e_defaults,
616 .pix_fmts = librav1e_pix_fmts,
617 .capabilities = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AUTO_THREADS,
618 .caps_internal = FF_CODEC_CAP_INIT_CLEANUP,
619 .wrapper_name = "librav1e",