MrBesen/ffmpeg
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity
1c7e55dd50
ffmpeg/libavcodec/librav1e.c
James Almer 827d6fe73d avcodec/encode: restructure the core encoding code 
This commit follows the same logic as 061a0c14bb, but for the encode API: The
new public encoding API will no longer be a wrapper around the old deprecated
one, and the internal API used by the encoders now consists of a single
receive_packet() callback that pulls frames as required.

amf encoders adapted by James Almer
librav1e encoder adapted by James Almer
nvidia encoders adapted by James Almer
MediaFoundation encoders adapted by James Almer
vaapi encoders adapted by Linjie Fu
v4l2_m2m encoders adapted by Andriy Gelman

Signed-off-by: James Almer <jamrial@gmail.com>
2020-06-18 17:11:37 -03:00

621 lines
19 KiB
C
Raw Blame History

/*
* 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 "encode.h"
#include "internal.h"
typedef struct librav1eContext {
const AVClass *class;
RaContext *ctx;
AVFrame *frame;
RaFrame *rframe;
AVBSFContext *bsf;
uint8_t *pass_data;
size_t pass_pos;
int pass_size;
AVDictionary *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;
}
if (ctx->rframe) {
rav1e_frame_unref(ctx->rframe);
ctx->rframe = NULL;
}
av_frame_free(&ctx->frame);
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;
ctx->frame = av_frame_alloc();
if (!ctx->frame)
return AVERROR(ENOMEM);
cfg = rav1e_config_default();
if (!cfg) {
av_log(avctx, AV_LOG_ERROR, "Could not allocate rav1e config.\n");
return AVERROR_EXTERNAL;
}
/*
* Rav1e currently uses the time base given to it only for ratecontrol... where
* the inverse is taken and used as a framerate. So, do what we do in other wrappers
* and use the framerate if we can.
*/
if (avctx->framerate.num > 0 && avctx->framerate.den > 0) {
rav1e_config_set_time_base(cfg, (RaRational) {
avctx->framerate.den, avctx->framerate.num
});
} else {
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_PASS1 || avctx->flags & AV_CODEC_FLAG_PASS2) && !avctx->bit_rate) {
av_log(avctx, AV_LOG_ERROR, "A bitrate must be set to use two pass mode.\n");
ret = AVERROR_INVALIDDATA;
goto end;
}
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;
}
}
{
AVDictionaryEntry *en = NULL;
while ((en = av_dict_get(ctx->rav1e_opts, "", 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);
}
}
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_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
{
librav1eContext *ctx = avctx->priv_data;
RaFrame *rframe = ctx->rframe;
RaPacket *rpkt = NULL;
int ret;
if (!rframe) {
AVFrame *frame = ctx->frame;
ret = ff_encode_get_frame(avctx, frame);
if (ret < 0 && ret != AVERROR_EOF)
return ret;
if (frame->buf[0]) {
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");
av_frame_unref(frame);
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);
}
av_frame_unref(frame);
}
}
ret = rav1e_send_frame(ctx->ctx, rframe);
if (rframe)
if (ret == RA_ENCODER_STATUS_ENOUGH_DATA) {
ctx->rframe = rframe; /* Queue is full. Store the RaFrame to retry next call */
} else {
rav1e_frame_unref(rframe); /* No need to unref if flushing. */
ctx->rframe = NULL;
}
switch (ret) {
case RA_ENCODER_STATUS_SUCCESS:
case RA_ENCODER_STATUS_ENOUGH_DATA:
break;
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;
}
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:
goto retry;
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;
rav1e_packet_unref(rpkt);
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");
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");
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 = -1 }, -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_DICT, { 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,
.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",
};
Reference in New Issue View Git Blame Copy Permalink