diff options
author | Derek Buitenhuis <derek.buitenhuis@gmail.com> | 2019-05-29 21:28:24 +0300 |
---|---|---|
committer | Derek Buitenhuis <derek.buitenhuis@gmail.com> | 2019-11-10 16:55:39 +0300 |
commit | d8bf24459b694338de4ceb2a2e6d4d2949d6658d (patch) | |
tree | 1613cbdf8e695b2f7ef484484bedf79ae0f515e7 /libavcodec/librav1e.c | |
parent | 0f89a2293ea5f642a67700225d76948ed154418e (diff) |
avcodec: Add librav1e encoder
Port to the new send/receive API by: James Almer <jamrial@gmail.com>.
Signed-off-by: Derek Buitenhuis <derek.buitenhuis@gmail.com>
Diffstat (limited to 'libavcodec/librav1e.c')
-rw-r--r-- | libavcodec/librav1e.c | 594 |
1 files changed, 594 insertions, 0 deletions
diff --git a/libavcodec/librav1e.c b/libavcodec/librav1e.c new file mode 100644 index 0000000000..f65d784127 --- /dev/null +++ b/libavcodec/librav1e.c @@ -0,0 +1,594 @@ +/* + * librav1e encoder + * + * Copyright (c) 2019 Derek Buitenhuis + * + * This file is part of FFmpeg. + * + * FFmpeg is free software; you can redistribute it and/or + * modify it under the terms of the GNU Lesser General Public + * License as published by the Free Software Foundation; either + * version 2.1 of the License, or (at your option) any later version. + * + * FFmpeg is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * Lesser General Public License for more details. + * + * You should have received a copy of the GNU Lesser General Public + * License along with FFmpeg; if not, write to the Free Software + * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA + */ + +#include <rav1e.h> + +#include "libavutil/internal.h" +#include "libavutil/avassert.h" +#include "libavutil/base64.h" +#include "libavutil/common.h" +#include "libavutil/mathematics.h" +#include "libavutil/opt.h" +#include "libavutil/pixdesc.h" +#include "avcodec.h" +#include "internal.h" + +typedef struct librav1eContext { + const AVClass *class; + + RaContext *ctx; + AVBSFContext *bsf; + + uint8_t *pass_data; + size_t pass_pos; + int pass_size; + + char *rav1e_opts; + int quantizer; + int speed; + int tiles; + int tile_rows; + int tile_cols; +} librav1eContext; + +static inline RaPixelRange range_map(enum AVPixelFormat pix_fmt, enum AVColorRange range) +{ + switch (pix_fmt) { + case AV_PIX_FMT_YUVJ420P: + case AV_PIX_FMT_YUVJ422P: + case AV_PIX_FMT_YUVJ444P: + return RA_PIXEL_RANGE_FULL; + } + + switch (range) { + case AVCOL_RANGE_JPEG: + return RA_PIXEL_RANGE_FULL; + case AVCOL_RANGE_MPEG: + default: + return RA_PIXEL_RANGE_LIMITED; + } +} + +static inline RaChromaSampling pix_fmt_map(enum AVPixelFormat pix_fmt) +{ + switch (pix_fmt) { + case AV_PIX_FMT_YUV420P: + case AV_PIX_FMT_YUVJ420P: + case AV_PIX_FMT_YUV420P10: + case AV_PIX_FMT_YUV420P12: + return RA_CHROMA_SAMPLING_CS420; + case AV_PIX_FMT_YUV422P: + case AV_PIX_FMT_YUVJ422P: + case AV_PIX_FMT_YUV422P10: + case AV_PIX_FMT_YUV422P12: + return RA_CHROMA_SAMPLING_CS422; + case AV_PIX_FMT_YUV444P: + case AV_PIX_FMT_YUVJ444P: + case AV_PIX_FMT_YUV444P10: + case AV_PIX_FMT_YUV444P12: + return RA_CHROMA_SAMPLING_CS444; + default: + av_assert0(0); + } +} + +static inline RaChromaSamplePosition chroma_loc_map(enum AVChromaLocation chroma_loc) +{ + switch (chroma_loc) { + case AVCHROMA_LOC_LEFT: + return RA_CHROMA_SAMPLE_POSITION_VERTICAL; + case AVCHROMA_LOC_TOPLEFT: + return RA_CHROMA_SAMPLE_POSITION_COLOCATED; + default: + return RA_CHROMA_SAMPLE_POSITION_UNKNOWN; + } +} + +static int get_stats(AVCodecContext *avctx, int eos) +{ + librav1eContext *ctx = avctx->priv_data; + RaData* buf = rav1e_twopass_out(ctx->ctx); + if (!buf) + return 0; + + if (!eos) { + uint8_t *tmp = av_fast_realloc(ctx->pass_data, &ctx->pass_size, + ctx->pass_pos + buf->len); + if (!tmp) { + rav1e_data_unref(buf); + return AVERROR(ENOMEM); + } + + ctx->pass_data = tmp; + memcpy(ctx->pass_data + ctx->pass_pos, buf->data, buf->len); + ctx->pass_pos += buf->len; + } else { + size_t b64_size = AV_BASE64_SIZE(ctx->pass_pos); + + memcpy(ctx->pass_data, buf->data, buf->len); + + avctx->stats_out = av_malloc(b64_size); + if (!avctx->stats_out) { + rav1e_data_unref(buf); + return AVERROR(ENOMEM); + } + + av_base64_encode(avctx->stats_out, b64_size, ctx->pass_data, ctx->pass_pos); + + av_freep(&ctx->pass_data); + } + + rav1e_data_unref(buf); + + return 0; +} + +static int set_stats(AVCodecContext *avctx) +{ + librav1eContext *ctx = avctx->priv_data; + int ret = 1; + + while (ret > 0 && ctx->pass_size - ctx->pass_pos > 0) { + ret = rav1e_twopass_in(ctx->ctx, ctx->pass_data + ctx->pass_pos, ctx->pass_size); + if (ret < 0) + return AVERROR_EXTERNAL; + ctx->pass_pos += ret; + } + + return 0; +} + +static av_cold int librav1e_encode_close(AVCodecContext *avctx) +{ + librav1eContext *ctx = avctx->priv_data; + + if (ctx->ctx) { + rav1e_context_unref(ctx->ctx); + ctx->ctx = NULL; + } + + av_bsf_free(&ctx->bsf); + av_freep(&ctx->pass_data); + + return 0; +} + +static av_cold int librav1e_encode_init(AVCodecContext *avctx) +{ + librav1eContext *ctx = avctx->priv_data; + const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(avctx->pix_fmt); + RaConfig *cfg = NULL; + int rret; + int ret = 0; + + cfg = rav1e_config_default(); + if (!cfg) { + av_log(avctx, AV_LOG_ERROR, "Could not allocate rav1e config.\n"); + return AVERROR_EXTERNAL; + } + + rav1e_config_set_time_base(cfg, (RaRational) { + avctx->time_base.num * avctx->ticks_per_frame, + avctx->time_base.den + }); + + if (avctx->flags & AV_CODEC_FLAG_PASS2) { + if (!avctx->stats_in) { + av_log(avctx, AV_LOG_ERROR, "No stats file provided for second pass.\n"); + ret = AVERROR(EINVAL); + goto end; + } + + ctx->pass_size = (strlen(avctx->stats_in) * 3) / 4; + ctx->pass_data = av_malloc(ctx->pass_size); + if (!ctx->pass_data) { + av_log(avctx, AV_LOG_ERROR, "Could not allocate stats buffer.\n"); + ret = AVERROR(ENOMEM); + goto end; + } + + ctx->pass_size = av_base64_decode(ctx->pass_data, avctx->stats_in, ctx->pass_size); + if (ctx->pass_size < 0) { + av_log(avctx, AV_LOG_ERROR, "Invalid pass file.\n"); + ret = AVERROR(EINVAL); + goto end; + } + } + + if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) { + const AVBitStreamFilter *filter = av_bsf_get_by_name("extract_extradata"); + int bret; + + if (!filter) { + av_log(avctx, AV_LOG_ERROR, "extract_extradata bitstream filter " + "not found. This is a bug, please report it.\n"); + ret = AVERROR_BUG; + goto end; + } + + bret = av_bsf_alloc(filter, &ctx->bsf); + if (bret < 0) { + ret = bret; + goto end; + } + + bret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx); + if (bret < 0) { + ret = bret; + goto end; + } + + bret = av_bsf_init(ctx->bsf); + if (bret < 0) { + ret = bret; + goto end; + } + } + + if (ctx->rav1e_opts) { + AVDictionary *dict = NULL; + AVDictionaryEntry *en = NULL; + + if (!av_dict_parse_string(&dict, ctx->rav1e_opts, "=", ":", 0)) { + while (en = av_dict_get(dict, "", en, AV_DICT_IGNORE_SUFFIX)) { + int parse_ret = rav1e_config_parse(cfg, en->key, en->value); + if (parse_ret < 0) + av_log(avctx, AV_LOG_WARNING, "Invalid value for %s: %s.\n", en->key, en->value); + } + av_dict_free(&dict); + } + } + + rret = rav1e_config_parse_int(cfg, "width", avctx->width); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Invalid width passed to rav1e.\n"); + ret = AVERROR_INVALIDDATA; + goto end; + } + + rret = rav1e_config_parse_int(cfg, "height", avctx->height); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Invalid height passed to rav1e.\n"); + ret = AVERROR_INVALIDDATA; + goto end; + } + + rret = rav1e_config_parse_int(cfg, "threads", avctx->thread_count); + if (rret < 0) + av_log(avctx, AV_LOG_WARNING, "Invalid number of threads, defaulting to auto.\n"); + + if (ctx->speed >= 0) { + rret = rav1e_config_parse_int(cfg, "speed", ctx->speed); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set speed preset.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + + /* rav1e handles precedence between 'tiles' and cols/rows for us. */ + if (ctx->tiles > 0) { + rret = rav1e_config_parse_int(cfg, "tiles", ctx->tiles); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set number of tiles to encode with.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + if (ctx->tile_rows > 0) { + rret = rav1e_config_parse_int(cfg, "tile_rows", ctx->tile_rows); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set number of tile rows to encode with.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + if (ctx->tile_cols > 0) { + rret = rav1e_config_parse_int(cfg, "tile_cols", ctx->tile_cols); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set number of tile cols to encode with.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + + if (avctx->gop_size > 0) { + rret = rav1e_config_parse_int(cfg, "key_frame_interval", avctx->gop_size); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set max keyint.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + + if (avctx->keyint_min > 0) { + rret = rav1e_config_parse_int(cfg, "min_key_frame_interval", avctx->keyint_min); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set min keyint.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + + if (avctx->bit_rate && ctx->quantizer < 0) { + int max_quantizer = avctx->qmax >= 0 ? avctx->qmax : 255; + + rret = rav1e_config_parse_int(cfg, "quantizer", max_quantizer); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set max quantizer.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + + if (avctx->qmin >= 0) { + rret = rav1e_config_parse_int(cfg, "min_quantizer", avctx->qmin); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set min quantizer.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + + rret = rav1e_config_parse_int(cfg, "bitrate", avctx->bit_rate); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set bitrate.\n"); + ret = AVERROR_INVALIDDATA; + goto end; + } + } else if (ctx->quantizer >= 0) { + if (avctx->bit_rate) + av_log(avctx, AV_LOG_WARNING, "Both bitrate and quantizer specified. Using quantizer mode."); + + rret = rav1e_config_parse_int(cfg, "quantizer", ctx->quantizer); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not set quantizer.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + } + + rret = rav1e_config_set_pixel_format(cfg, desc->comp[0].depth, + pix_fmt_map(avctx->pix_fmt), + chroma_loc_map(avctx->chroma_sample_location), + range_map(avctx->pix_fmt, avctx->color_range)); + if (rret < 0) { + av_log(avctx, AV_LOG_ERROR, "Failed to set pixel format properties.\n"); + ret = AVERROR_INVALIDDATA; + goto end; + } + + /* rav1e's colorspace enums match standard values. */ + rret = rav1e_config_set_color_description(cfg, (RaMatrixCoefficients) avctx->colorspace, + (RaColorPrimaries) avctx->color_primaries, + (RaTransferCharacteristics) avctx->color_trc); + if (rret < 0) { + av_log(avctx, AV_LOG_WARNING, "Failed to set color properties.\n"); + if (avctx->err_recognition & AV_EF_EXPLODE) { + ret = AVERROR_INVALIDDATA; + goto end; + } + } + + ctx->ctx = rav1e_context_new(cfg); + if (!ctx->ctx) { + av_log(avctx, AV_LOG_ERROR, "Failed to create rav1e encode context.\n"); + ret = AVERROR_EXTERNAL; + goto end; + } + + ret = 0; + +end: + + rav1e_config_unref(cfg); + + return ret; +} + +static int librav1e_send_frame(AVCodecContext *avctx, const AVFrame *frame) +{ + librav1eContext *ctx = avctx->priv_data; + RaFrame *rframe = NULL; + int ret; + + if (frame) { + const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(frame->format); + + rframe = rav1e_frame_new(ctx->ctx); + if (!rframe) { + av_log(avctx, AV_LOG_ERROR, "Could not allocate new rav1e frame.\n"); + return AVERROR(ENOMEM); + } + + for (int i = 0; i < desc->nb_components; i++) { + int shift = i ? desc->log2_chroma_h : 0; + int bytes = desc->comp[0].depth == 8 ? 1 : 2; + rav1e_frame_fill_plane(rframe, i, frame->data[i], + (frame->height >> shift) * frame->linesize[i], + frame->linesize[i], bytes); + } + } + + ret = rav1e_send_frame(ctx->ctx, rframe); + if (rframe) + rav1e_frame_unref(rframe); /* No need to unref if flushing. */ + + switch (ret) { + case RA_ENCODER_STATUS_SUCCESS: + break; + case RA_ENCODER_STATUS_ENOUGH_DATA: + return AVERROR(EAGAIN); + case RA_ENCODER_STATUS_FAILURE: + av_log(avctx, AV_LOG_ERROR, "Could not send frame: %s\n", rav1e_status_to_str(ret)); + return AVERROR_EXTERNAL; + default: + av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_send_frame: %s\n", ret, rav1e_status_to_str(ret)); + return AVERROR_UNKNOWN; + } + + return 0; +} + +static int librav1e_receive_packet(AVCodecContext *avctx, AVPacket *pkt) +{ + librav1eContext *ctx = avctx->priv_data; + RaPacket *rpkt = NULL; + int ret; + +retry: + + if (avctx->flags & AV_CODEC_FLAG_PASS1) { + int sret = get_stats(avctx, 0); + if (sret < 0) + return sret; + } else if (avctx->flags & AV_CODEC_FLAG_PASS2) { + int sret = set_stats(avctx); + if (sret < 0) + return sret; + } + + ret = rav1e_receive_packet(ctx->ctx, &rpkt); + switch (ret) { + case RA_ENCODER_STATUS_SUCCESS: + break; + case RA_ENCODER_STATUS_LIMIT_REACHED: + if (avctx->flags & AV_CODEC_FLAG_PASS1) { + int sret = get_stats(avctx, 1); + if (sret < 0) + return sret; + } + return AVERROR_EOF; + case RA_ENCODER_STATUS_ENCODED: + if (avctx->internal->draining) + goto retry; + return AVERROR(EAGAIN); + case RA_ENCODER_STATUS_NEED_MORE_DATA: + if (avctx->internal->draining) { + av_log(avctx, AV_LOG_ERROR, "Unexpected error when receiving packet after EOF.\n"); + return AVERROR_EXTERNAL; + } + return AVERROR(EAGAIN); + case RA_ENCODER_STATUS_FAILURE: + av_log(avctx, AV_LOG_ERROR, "Could not encode frame: %s\n", rav1e_status_to_str(ret)); + return AVERROR_EXTERNAL; + default: + av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_receive_packet: %s\n", ret, rav1e_status_to_str(ret)); + return AVERROR_UNKNOWN; + } + + ret = av_new_packet(pkt, rpkt->len); + if (ret < 0) { + av_log(avctx, AV_LOG_ERROR, "Could not allocate packet.\n"); + rav1e_packet_unref(rpkt); + return ret; + } + + memcpy(pkt->data, rpkt->data, rpkt->len); + + if (rpkt->frame_type == RA_FRAME_TYPE_KEY) + pkt->flags |= AV_PKT_FLAG_KEY; + + pkt->pts = pkt->dts = rpkt->input_frameno * avctx->ticks_per_frame; + + if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) { + int ret = av_bsf_send_packet(ctx->bsf, pkt); + if (ret < 0) { + av_log(avctx, AV_LOG_ERROR, "extradata extraction send failed.\n"); + rav1e_packet_unref(rpkt); + av_packet_unref(pkt); + return ret; + } + + ret = av_bsf_receive_packet(ctx->bsf, pkt); + if (ret < 0) { + av_log(avctx, AV_LOG_ERROR, "extradata extraction receive failed.\n"); + rav1e_packet_unref(rpkt); + av_packet_unref(pkt); + return ret; + } + } + + return 0; +} + +#define OFFSET(x) offsetof(librav1eContext, x) +#define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM + +static const AVOption options[] = { + { "qp", "use constant quantizer mode", OFFSET(quantizer), AV_OPT_TYPE_INT, { .i64 = 100 }, -1, 255, VE }, + { "speed", "what speed preset to use", OFFSET(speed), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 10, VE }, + { "tiles", "number of tiles encode with", OFFSET(tiles), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE }, + { "tile-rows", "number of tiles rows to encode with", OFFSET(tile_rows), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE }, + { "tile-columns", "number of tiles columns to encode with", OFFSET(tile_cols), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE }, + { "rav1e-params", "set the rav1e configuration using a :-separated list of key=value parameters", OFFSET(rav1e_opts), AV_OPT_TYPE_STRING, { 0 }, 0, 0, VE }, + { NULL } +}; + +static const AVCodecDefault librav1e_defaults[] = { + { "b", "0" }, + { "g", "0" }, + { "keyint_min", "0" }, + { "qmax", "-1" }, + { "qmin", "-1" }, + { NULL } +}; + +const enum AVPixelFormat librav1e_pix_fmts[] = { + AV_PIX_FMT_YUV420P, + AV_PIX_FMT_YUVJ420P, + AV_PIX_FMT_YUV420P10, + AV_PIX_FMT_YUV420P12, + AV_PIX_FMT_YUV422P, + AV_PIX_FMT_YUVJ422P, + AV_PIX_FMT_YUV422P10, + AV_PIX_FMT_YUV422P12, + AV_PIX_FMT_YUV444P, + AV_PIX_FMT_YUVJ444P, + AV_PIX_FMT_YUV444P10, + AV_PIX_FMT_YUV444P12, + AV_PIX_FMT_NONE +}; + +static const AVClass class = { + .class_name = "librav1e", + .item_name = av_default_item_name, + .option = options, + .version = LIBAVUTIL_VERSION_INT, +}; + +AVCodec ff_librav1e_encoder = { + .name = "librav1e", + .long_name = NULL_IF_CONFIG_SMALL("librav1e AV1"), + .type = AVMEDIA_TYPE_VIDEO, + .id = AV_CODEC_ID_AV1, + .init = librav1e_encode_init, + .send_frame = librav1e_send_frame, + .receive_packet = librav1e_receive_packet, + .close = librav1e_encode_close, + .priv_data_size = sizeof(librav1eContext), + .priv_class = &class, + .defaults = librav1e_defaults, + .pix_fmts = librav1e_pix_fmts, + .capabilities = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AUTO_THREADS, + .caps_internal = FF_CODEC_CAP_INIT_CLEANUP, + .wrapper_name = "librav1e", +}; |