]> git.sesse.net Git - ffmpeg/blob - libavcodec/librav1e.c
avcodec/librav1e: Use the framerate when available for ratecontrol
[ffmpeg] / libavcodec / librav1e.c
1 /*
2  * librav1e encoder
3  *
4  * Copyright (c) 2019 Derek Buitenhuis
5  *
6  * This file is part of FFmpeg.
7  *
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.
12  *
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.
17  *
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
21  */
22
23 #include <rav1e.h>
24
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"
32 #include "avcodec.h"
33 #include "internal.h"
34
35 typedef struct librav1eContext {
36     const AVClass *class;
37
38     RaContext *ctx;
39     AVBSFContext *bsf;
40
41     uint8_t *pass_data;
42     size_t pass_pos;
43     int pass_size;
44
45     AVDictionary *rav1e_opts;
46     int quantizer;
47     int speed;
48     int tiles;
49     int tile_rows;
50     int tile_cols;
51 } librav1eContext;
52
53 static inline RaPixelRange range_map(enum AVPixelFormat pix_fmt, enum AVColorRange range)
54 {
55     switch (pix_fmt) {
56     case AV_PIX_FMT_YUVJ420P:
57     case AV_PIX_FMT_YUVJ422P:
58     case AV_PIX_FMT_YUVJ444P:
59         return RA_PIXEL_RANGE_FULL;
60     }
61
62     switch (range) {
63     case AVCOL_RANGE_JPEG:
64         return RA_PIXEL_RANGE_FULL;
65     case AVCOL_RANGE_MPEG:
66     default:
67         return RA_PIXEL_RANGE_LIMITED;
68     }
69 }
70
71 static inline RaChromaSampling pix_fmt_map(enum AVPixelFormat pix_fmt)
72 {
73     switch (pix_fmt) {
74     case AV_PIX_FMT_YUV420P:
75     case AV_PIX_FMT_YUVJ420P:
76     case AV_PIX_FMT_YUV420P10:
77     case AV_PIX_FMT_YUV420P12:
78         return RA_CHROMA_SAMPLING_CS420;
79     case AV_PIX_FMT_YUV422P:
80     case AV_PIX_FMT_YUVJ422P:
81     case AV_PIX_FMT_YUV422P10:
82     case AV_PIX_FMT_YUV422P12:
83         return RA_CHROMA_SAMPLING_CS422;
84     case AV_PIX_FMT_YUV444P:
85     case AV_PIX_FMT_YUVJ444P:
86     case AV_PIX_FMT_YUV444P10:
87     case AV_PIX_FMT_YUV444P12:
88         return RA_CHROMA_SAMPLING_CS444;
89     default:
90         av_assert0(0);
91     }
92 }
93
94 static inline RaChromaSamplePosition chroma_loc_map(enum AVChromaLocation chroma_loc)
95 {
96     switch (chroma_loc) {
97     case AVCHROMA_LOC_LEFT:
98         return RA_CHROMA_SAMPLE_POSITION_VERTICAL;
99     case AVCHROMA_LOC_TOPLEFT:
100         return RA_CHROMA_SAMPLE_POSITION_COLOCATED;
101     default:
102         return RA_CHROMA_SAMPLE_POSITION_UNKNOWN;
103     }
104 }
105
106 static int get_stats(AVCodecContext *avctx, int eos)
107 {
108     librav1eContext *ctx = avctx->priv_data;
109     RaData* buf = rav1e_twopass_out(ctx->ctx);
110     if (!buf)
111         return 0;
112
113     if (!eos) {
114         uint8_t *tmp = av_fast_realloc(ctx->pass_data, &ctx->pass_size,
115                                       ctx->pass_pos + buf->len);
116         if (!tmp) {
117             rav1e_data_unref(buf);
118             return AVERROR(ENOMEM);
119         }
120
121         ctx->pass_data = tmp;
122         memcpy(ctx->pass_data + ctx->pass_pos, buf->data, buf->len);
123         ctx->pass_pos += buf->len;
124     } else {
125         size_t b64_size = AV_BASE64_SIZE(ctx->pass_pos);
126
127         memcpy(ctx->pass_data, buf->data, buf->len);
128
129         avctx->stats_out = av_malloc(b64_size);
130         if (!avctx->stats_out) {
131             rav1e_data_unref(buf);
132             return AVERROR(ENOMEM);
133         }
134
135         av_base64_encode(avctx->stats_out, b64_size, ctx->pass_data, ctx->pass_pos);
136
137         av_freep(&ctx->pass_data);
138     }
139
140     rav1e_data_unref(buf);
141
142     return 0;
143 }
144
145 static int set_stats(AVCodecContext *avctx)
146 {
147     librav1eContext *ctx = avctx->priv_data;
148     int ret = 1;
149
150     while (ret > 0 && ctx->pass_size - ctx->pass_pos > 0) {
151         ret = rav1e_twopass_in(ctx->ctx, ctx->pass_data + ctx->pass_pos, ctx->pass_size);
152         if (ret < 0)
153             return AVERROR_EXTERNAL;
154         ctx->pass_pos += ret;
155     }
156
157     return 0;
158 }
159
160 static av_cold int librav1e_encode_close(AVCodecContext *avctx)
161 {
162     librav1eContext *ctx = avctx->priv_data;
163
164     if (ctx->ctx) {
165         rav1e_context_unref(ctx->ctx);
166         ctx->ctx = NULL;
167     }
168
169     av_bsf_free(&ctx->bsf);
170     av_freep(&ctx->pass_data);
171
172     return 0;
173 }
174
175 static av_cold int librav1e_encode_init(AVCodecContext *avctx)
176 {
177     librav1eContext *ctx = avctx->priv_data;
178     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(avctx->pix_fmt);
179     RaConfig *cfg = NULL;
180     int rret;
181     int ret = 0;
182
183     cfg = rav1e_config_default();
184     if (!cfg) {
185         av_log(avctx, AV_LOG_ERROR, "Could not allocate rav1e config.\n");
186         return AVERROR_EXTERNAL;
187     }
188
189     /*
190      * Rav1e currently uses the time base given to it only for ratecontrol... where
191      * the inverse is taken and used as a framerate. So, do what we do in other wrappers
192      * and use the framerate if we can.
193      */
194     if (avctx->framerate.num > 0 && avctx->framerate.den > 0) {
195         rav1e_config_set_time_base(cfg, (RaRational) {
196                                    avctx->framerate.den, avctx->framerate.num
197                                    });
198     } else {
199         rav1e_config_set_time_base(cfg, (RaRational) {
200                                    avctx->time_base.num * avctx->ticks_per_frame,
201                                    avctx->time_base.den
202                                    });
203     }
204
205     if (avctx->flags & AV_CODEC_FLAG_PASS2) {
206         if (!avctx->stats_in) {
207             av_log(avctx, AV_LOG_ERROR, "No stats file provided for second pass.\n");
208             ret = AVERROR(EINVAL);
209             goto end;
210         }
211
212         ctx->pass_size = (strlen(avctx->stats_in) * 3) / 4;
213         ctx->pass_data = av_malloc(ctx->pass_size);
214         if (!ctx->pass_data) {
215             av_log(avctx, AV_LOG_ERROR, "Could not allocate stats buffer.\n");
216             ret = AVERROR(ENOMEM);
217             goto end;
218         }
219
220         ctx->pass_size = av_base64_decode(ctx->pass_data, avctx->stats_in, ctx->pass_size);
221         if (ctx->pass_size < 0) {
222             av_log(avctx, AV_LOG_ERROR, "Invalid pass file.\n");
223             ret = AVERROR(EINVAL);
224             goto end;
225         }
226     }
227
228     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
229          const AVBitStreamFilter *filter = av_bsf_get_by_name("extract_extradata");
230          int bret;
231
232          if (!filter) {
233             av_log(avctx, AV_LOG_ERROR, "extract_extradata bitstream filter "
234                    "not found. This is a bug, please report it.\n");
235             ret = AVERROR_BUG;
236             goto end;
237          }
238
239          bret = av_bsf_alloc(filter, &ctx->bsf);
240          if (bret < 0) {
241              ret = bret;
242              goto end;
243          }
244
245          bret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx);
246          if (bret < 0) {
247              ret = bret;
248              goto end;
249          }
250
251          bret = av_bsf_init(ctx->bsf);
252          if (bret < 0) {
253              ret = bret;
254              goto end;
255          }
256     }
257
258     {
259         AVDictionaryEntry *en = NULL;
260         while ((en = av_dict_get(ctx->rav1e_opts, "", en, AV_DICT_IGNORE_SUFFIX))) {
261             int parse_ret = rav1e_config_parse(cfg, en->key, en->value);
262             if (parse_ret < 0)
263                 av_log(avctx, AV_LOG_WARNING, "Invalid value for %s: %s.\n", en->key, en->value);
264         }
265     }
266
267     rret = rav1e_config_parse_int(cfg, "width", avctx->width);
268     if (rret < 0) {
269         av_log(avctx, AV_LOG_ERROR, "Invalid width passed to rav1e.\n");
270         ret = AVERROR_INVALIDDATA;
271         goto end;
272     }
273
274     rret = rav1e_config_parse_int(cfg, "height", avctx->height);
275     if (rret < 0) {
276         av_log(avctx, AV_LOG_ERROR, "Invalid height passed to rav1e.\n");
277         ret = AVERROR_INVALIDDATA;
278         goto end;
279     }
280
281     rret = rav1e_config_parse_int(cfg, "threads", avctx->thread_count);
282     if (rret < 0)
283         av_log(avctx, AV_LOG_WARNING, "Invalid number of threads, defaulting to auto.\n");
284
285     if (ctx->speed >= 0) {
286         rret = rav1e_config_parse_int(cfg, "speed", ctx->speed);
287         if (rret < 0) {
288             av_log(avctx, AV_LOG_ERROR, "Could not set speed preset.\n");
289             ret = AVERROR_EXTERNAL;
290             goto end;
291         }
292     }
293
294     /* rav1e handles precedence between 'tiles' and cols/rows for us. */
295     if (ctx->tiles > 0) {
296         rret = rav1e_config_parse_int(cfg, "tiles", ctx->tiles);
297         if (rret < 0) {
298             av_log(avctx, AV_LOG_ERROR, "Could not set number of tiles to encode with.\n");
299             ret = AVERROR_EXTERNAL;
300             goto end;
301         }
302     }
303     if (ctx->tile_rows > 0) {
304         rret = rav1e_config_parse_int(cfg, "tile_rows", ctx->tile_rows);
305         if (rret < 0) {
306             av_log(avctx, AV_LOG_ERROR, "Could not set number of tile rows to encode with.\n");
307             ret = AVERROR_EXTERNAL;
308             goto end;
309         }
310     }
311     if (ctx->tile_cols > 0) {
312         rret = rav1e_config_parse_int(cfg, "tile_cols", ctx->tile_cols);
313         if (rret < 0) {
314             av_log(avctx, AV_LOG_ERROR, "Could not set number of tile cols to encode with.\n");
315             ret = AVERROR_EXTERNAL;
316             goto end;
317         }
318     }
319
320     if (avctx->gop_size > 0) {
321         rret = rav1e_config_parse_int(cfg, "key_frame_interval", avctx->gop_size);
322         if (rret < 0) {
323             av_log(avctx, AV_LOG_ERROR, "Could not set max keyint.\n");
324             ret = AVERROR_EXTERNAL;
325             goto end;
326         }
327     }
328
329     if (avctx->keyint_min > 0) {
330         rret = rav1e_config_parse_int(cfg, "min_key_frame_interval", avctx->keyint_min);
331         if (rret < 0) {
332             av_log(avctx, AV_LOG_ERROR, "Could not set min keyint.\n");
333             ret = AVERROR_EXTERNAL;
334             goto end;
335         }
336     }
337
338     if (avctx->bit_rate && ctx->quantizer < 0) {
339         int max_quantizer = avctx->qmax >= 0 ? avctx->qmax : 255;
340
341         rret = rav1e_config_parse_int(cfg, "quantizer", max_quantizer);
342         if (rret < 0) {
343             av_log(avctx, AV_LOG_ERROR, "Could not set max quantizer.\n");
344             ret = AVERROR_EXTERNAL;
345             goto end;
346         }
347
348         if (avctx->qmin >= 0) {
349             rret = rav1e_config_parse_int(cfg, "min_quantizer", avctx->qmin);
350             if (rret < 0) {
351                 av_log(avctx, AV_LOG_ERROR, "Could not set min quantizer.\n");
352                 ret = AVERROR_EXTERNAL;
353                 goto end;
354             }
355         }
356
357         rret = rav1e_config_parse_int(cfg, "bitrate", avctx->bit_rate);
358         if (rret < 0) {
359             av_log(avctx, AV_LOG_ERROR, "Could not set bitrate.\n");
360             ret = AVERROR_INVALIDDATA;
361             goto end;
362         }
363     } else if (ctx->quantizer >= 0) {
364         if (avctx->bit_rate)
365             av_log(avctx, AV_LOG_WARNING, "Both bitrate and quantizer specified. Using quantizer mode.");
366
367         rret = rav1e_config_parse_int(cfg, "quantizer", ctx->quantizer);
368         if (rret < 0) {
369             av_log(avctx, AV_LOG_ERROR, "Could not set quantizer.\n");
370             ret = AVERROR_EXTERNAL;
371             goto end;
372         }
373     }
374
375     rret = rav1e_config_set_pixel_format(cfg, desc->comp[0].depth,
376                                          pix_fmt_map(avctx->pix_fmt),
377                                          chroma_loc_map(avctx->chroma_sample_location),
378                                          range_map(avctx->pix_fmt, avctx->color_range));
379     if (rret < 0) {
380         av_log(avctx, AV_LOG_ERROR, "Failed to set pixel format properties.\n");
381         ret = AVERROR_INVALIDDATA;
382         goto end;
383     }
384
385     /* rav1e's colorspace enums match standard values. */
386     rret = rav1e_config_set_color_description(cfg, (RaMatrixCoefficients) avctx->colorspace,
387                                               (RaColorPrimaries) avctx->color_primaries,
388                                               (RaTransferCharacteristics) avctx->color_trc);
389     if (rret < 0) {
390         av_log(avctx, AV_LOG_WARNING, "Failed to set color properties.\n");
391         if (avctx->err_recognition & AV_EF_EXPLODE) {
392             ret = AVERROR_INVALIDDATA;
393             goto end;
394         }
395     }
396
397     ctx->ctx = rav1e_context_new(cfg);
398     if (!ctx->ctx) {
399         av_log(avctx, AV_LOG_ERROR, "Failed to create rav1e encode context.\n");
400         ret = AVERROR_EXTERNAL;
401         goto end;
402     }
403
404     ret = 0;
405
406 end:
407
408     rav1e_config_unref(cfg);
409
410     return ret;
411 }
412
413 static int librav1e_send_frame(AVCodecContext *avctx, const AVFrame *frame)
414 {
415     librav1eContext *ctx = avctx->priv_data;
416     RaFrame *rframe = NULL;
417     int ret;
418
419     if (frame) {
420         const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(frame->format);
421
422         rframe = rav1e_frame_new(ctx->ctx);
423         if (!rframe) {
424             av_log(avctx, AV_LOG_ERROR, "Could not allocate new rav1e frame.\n");
425             return AVERROR(ENOMEM);
426         }
427
428         for (int i = 0; i < desc->nb_components; i++) {
429             int shift = i ? desc->log2_chroma_h : 0;
430             int bytes = desc->comp[0].depth == 8 ? 1 : 2;
431             rav1e_frame_fill_plane(rframe, i, frame->data[i],
432                                    (frame->height >> shift) * frame->linesize[i],
433                                    frame->linesize[i], bytes);
434         }
435     }
436
437     ret = rav1e_send_frame(ctx->ctx, rframe);
438     if (rframe)
439          rav1e_frame_unref(rframe); /* No need to unref if flushing. */
440
441     switch (ret) {
442     case RA_ENCODER_STATUS_SUCCESS:
443         break;
444     case RA_ENCODER_STATUS_ENOUGH_DATA:
445         return AVERROR(EAGAIN);
446     case RA_ENCODER_STATUS_FAILURE:
447         av_log(avctx, AV_LOG_ERROR, "Could not send frame: %s\n", rav1e_status_to_str(ret));
448         return AVERROR_EXTERNAL;
449     default:
450         av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_send_frame: %s\n", ret, rav1e_status_to_str(ret));
451         return AVERROR_UNKNOWN;
452     }
453
454     return 0;
455 }
456
457 static int librav1e_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
458 {
459     librav1eContext *ctx = avctx->priv_data;
460     RaPacket *rpkt = NULL;
461     int ret;
462
463 retry:
464
465     if (avctx->flags & AV_CODEC_FLAG_PASS1) {
466         int sret = get_stats(avctx, 0);
467         if (sret < 0)
468             return sret;
469     } else if (avctx->flags & AV_CODEC_FLAG_PASS2) {
470         int sret = set_stats(avctx);
471         if (sret < 0)
472             return sret;
473     }
474
475     ret = rav1e_receive_packet(ctx->ctx, &rpkt);
476     switch (ret) {
477     case RA_ENCODER_STATUS_SUCCESS:
478         break;
479     case RA_ENCODER_STATUS_LIMIT_REACHED:
480         if (avctx->flags & AV_CODEC_FLAG_PASS1) {
481             int sret = get_stats(avctx, 1);
482             if (sret < 0)
483                 return sret;
484         }
485         return AVERROR_EOF;
486     case RA_ENCODER_STATUS_ENCODED:
487         if (avctx->internal->draining)
488             goto retry;
489         return AVERROR(EAGAIN);
490     case RA_ENCODER_STATUS_NEED_MORE_DATA:
491         if (avctx->internal->draining) {
492             av_log(avctx, AV_LOG_ERROR, "Unexpected error when receiving packet after EOF.\n");
493             return AVERROR_EXTERNAL;
494         }
495         return AVERROR(EAGAIN);
496     case RA_ENCODER_STATUS_FAILURE:
497         av_log(avctx, AV_LOG_ERROR, "Could not encode frame: %s\n", rav1e_status_to_str(ret));
498         return AVERROR_EXTERNAL;
499     default:
500         av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_receive_packet: %s\n", ret, rav1e_status_to_str(ret));
501         return AVERROR_UNKNOWN;
502     }
503
504     ret = av_new_packet(pkt, rpkt->len);
505     if (ret < 0) {
506         av_log(avctx, AV_LOG_ERROR, "Could not allocate packet.\n");
507         rav1e_packet_unref(rpkt);
508         return ret;
509     }
510
511     memcpy(pkt->data, rpkt->data, rpkt->len);
512
513     if (rpkt->frame_type == RA_FRAME_TYPE_KEY)
514         pkt->flags |= AV_PKT_FLAG_KEY;
515
516     pkt->pts = pkt->dts = rpkt->input_frameno * avctx->ticks_per_frame;
517     rav1e_packet_unref(rpkt);
518
519     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
520         int ret = av_bsf_send_packet(ctx->bsf, pkt);
521         if (ret < 0) {
522             av_log(avctx, AV_LOG_ERROR, "extradata extraction send failed.\n");
523             av_packet_unref(pkt);
524             return ret;
525         }
526
527         ret = av_bsf_receive_packet(ctx->bsf, pkt);
528         if (ret < 0) {
529             av_log(avctx, AV_LOG_ERROR, "extradata extraction receive failed.\n");
530             av_packet_unref(pkt);
531             return ret;
532         }
533     }
534
535     return 0;
536 }
537
538 #define OFFSET(x) offsetof(librav1eContext, x)
539 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
540
541 static const AVOption options[] = {
542     { "qp", "use constant quantizer mode", OFFSET(quantizer), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 255, VE },
543     { "speed", "what speed preset to use", OFFSET(speed), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 10, VE },
544     { "tiles", "number of tiles encode with", OFFSET(tiles), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
545     { "tile-rows", "number of tiles rows to encode with", OFFSET(tile_rows), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
546     { "tile-columns", "number of tiles columns to encode with", OFFSET(tile_cols), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
547     { "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 },
548     { NULL }
549 };
550
551 static const AVCodecDefault librav1e_defaults[] = {
552     { "b",           "0" },
553     { "g",           "0" },
554     { "keyint_min",  "0" },
555     { "qmax",       "-1" },
556     { "qmin",       "-1" },
557     { NULL }
558 };
559
560 const enum AVPixelFormat librav1e_pix_fmts[] = {
561     AV_PIX_FMT_YUV420P,
562     AV_PIX_FMT_YUVJ420P,
563     AV_PIX_FMT_YUV420P10,
564     AV_PIX_FMT_YUV420P12,
565     AV_PIX_FMT_YUV422P,
566     AV_PIX_FMT_YUVJ422P,
567     AV_PIX_FMT_YUV422P10,
568     AV_PIX_FMT_YUV422P12,
569     AV_PIX_FMT_YUV444P,
570     AV_PIX_FMT_YUVJ444P,
571     AV_PIX_FMT_YUV444P10,
572     AV_PIX_FMT_YUV444P12,
573     AV_PIX_FMT_NONE
574 };
575
576 static const AVClass class = {
577     .class_name = "librav1e",
578     .item_name  = av_default_item_name,
579     .option     = options,
580     .version    = LIBAVUTIL_VERSION_INT,
581 };
582
583 AVCodec ff_librav1e_encoder = {
584     .name           = "librav1e",
585     .long_name      = NULL_IF_CONFIG_SMALL("librav1e AV1"),
586     .type           = AVMEDIA_TYPE_VIDEO,
587     .id             = AV_CODEC_ID_AV1,
588     .init           = librav1e_encode_init,
589     .send_frame     = librav1e_send_frame,
590     .receive_packet = librav1e_receive_packet,
591     .close          = librav1e_encode_close,
592     .priv_data_size = sizeof(librav1eContext),
593     .priv_class     = &class,
594     .defaults       = librav1e_defaults,
595     .pix_fmts       = librav1e_pix_fmts,
596     .capabilities   = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AUTO_THREADS,
597     .caps_internal  = FF_CODEC_CAP_INIT_CLEANUP,
598     .wrapper_name   = "librav1e",
599 };