avcodec: Add librav1e encoder
[ffmpeg.git] / 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     char *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     rav1e_config_set_time_base(cfg, (RaRational) {
190                                avctx->time_base.num * avctx->ticks_per_frame,
191                                avctx->time_base.den
192                                });
193
194     if (avctx->flags & AV_CODEC_FLAG_PASS2) {
195         if (!avctx->stats_in) {
196             av_log(avctx, AV_LOG_ERROR, "No stats file provided for second pass.\n");
197             ret = AVERROR(EINVAL);
198             goto end;
199         }
200
201         ctx->pass_size = (strlen(avctx->stats_in) * 3) / 4;
202         ctx->pass_data = av_malloc(ctx->pass_size);
203         if (!ctx->pass_data) {
204             av_log(avctx, AV_LOG_ERROR, "Could not allocate stats buffer.\n");
205             ret = AVERROR(ENOMEM);
206             goto end;
207         }
208
209         ctx->pass_size = av_base64_decode(ctx->pass_data, avctx->stats_in, ctx->pass_size);
210         if (ctx->pass_size < 0) {
211             av_log(avctx, AV_LOG_ERROR, "Invalid pass file.\n");
212             ret = AVERROR(EINVAL);
213             goto end;
214         }
215     }
216
217     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
218          const AVBitStreamFilter *filter = av_bsf_get_by_name("extract_extradata");
219          int bret;
220
221          if (!filter) {
222             av_log(avctx, AV_LOG_ERROR, "extract_extradata bitstream filter "
223                    "not found. This is a bug, please report it.\n");
224             ret = AVERROR_BUG;
225             goto end;
226          }
227
228          bret = av_bsf_alloc(filter, &ctx->bsf);
229          if (bret < 0) {
230              ret = bret;
231              goto end;
232          }
233
234          bret = avcodec_parameters_from_context(ctx->bsf->par_in, avctx);
235          if (bret < 0) {
236              ret = bret;
237              goto end;
238          }
239
240          bret = av_bsf_init(ctx->bsf);
241          if (bret < 0) {
242              ret = bret;
243              goto end;
244          }
245     }
246
247     if (ctx->rav1e_opts) {
248         AVDictionary *dict    = NULL;
249         AVDictionaryEntry *en = NULL;
250
251         if (!av_dict_parse_string(&dict, ctx->rav1e_opts, "=", ":", 0)) {
252             while (en = av_dict_get(dict, "", en, AV_DICT_IGNORE_SUFFIX)) {
253                 int parse_ret = rav1e_config_parse(cfg, en->key, en->value);
254                 if (parse_ret < 0)
255                     av_log(avctx, AV_LOG_WARNING, "Invalid value for %s: %s.\n", en->key, en->value);
256             }
257             av_dict_free(&dict);
258         }
259     }
260
261     rret = rav1e_config_parse_int(cfg, "width", avctx->width);
262     if (rret < 0) {
263         av_log(avctx, AV_LOG_ERROR, "Invalid width passed to rav1e.\n");
264         ret = AVERROR_INVALIDDATA;
265         goto end;
266     }
267
268     rret = rav1e_config_parse_int(cfg, "height", avctx->height);
269     if (rret < 0) {
270         av_log(avctx, AV_LOG_ERROR, "Invalid height passed to rav1e.\n");
271         ret = AVERROR_INVALIDDATA;
272         goto end;
273     }
274
275     rret = rav1e_config_parse_int(cfg, "threads", avctx->thread_count);
276     if (rret < 0)
277         av_log(avctx, AV_LOG_WARNING, "Invalid number of threads, defaulting to auto.\n");
278
279     if (ctx->speed >= 0) {
280         rret = rav1e_config_parse_int(cfg, "speed", ctx->speed);
281         if (rret < 0) {
282             av_log(avctx, AV_LOG_ERROR, "Could not set speed preset.\n");
283             ret = AVERROR_EXTERNAL;
284             goto end;
285         }
286     }
287
288     /* rav1e handles precedence between 'tiles' and cols/rows for us. */
289     if (ctx->tiles > 0) {
290         rret = rav1e_config_parse_int(cfg, "tiles", ctx->tiles);
291         if (rret < 0) {
292             av_log(avctx, AV_LOG_ERROR, "Could not set number of tiles to encode with.\n");
293             ret = AVERROR_EXTERNAL;
294             goto end;
295         }
296     }
297     if (ctx->tile_rows > 0) {
298         rret = rav1e_config_parse_int(cfg, "tile_rows", ctx->tile_rows);
299         if (rret < 0) {
300             av_log(avctx, AV_LOG_ERROR, "Could not set number of tile rows to encode with.\n");
301             ret = AVERROR_EXTERNAL;
302             goto end;
303         }
304     }
305     if (ctx->tile_cols > 0) {
306         rret = rav1e_config_parse_int(cfg, "tile_cols", ctx->tile_cols);
307         if (rret < 0) {
308             av_log(avctx, AV_LOG_ERROR, "Could not set number of tile cols to encode with.\n");
309             ret = AVERROR_EXTERNAL;
310             goto end;
311         }
312     }
313
314     if (avctx->gop_size > 0) {
315         rret = rav1e_config_parse_int(cfg, "key_frame_interval", avctx->gop_size);
316         if (rret < 0) {
317             av_log(avctx, AV_LOG_ERROR, "Could not set max keyint.\n");
318             ret = AVERROR_EXTERNAL;
319             goto end;
320         }
321     }
322
323     if (avctx->keyint_min > 0) {
324         rret = rav1e_config_parse_int(cfg, "min_key_frame_interval", avctx->keyint_min);
325         if (rret < 0) {
326             av_log(avctx, AV_LOG_ERROR, "Could not set min keyint.\n");
327             ret = AVERROR_EXTERNAL;
328             goto end;
329         }
330     }
331
332     if (avctx->bit_rate && ctx->quantizer < 0) {
333         int max_quantizer = avctx->qmax >= 0 ? avctx->qmax : 255;
334
335         rret = rav1e_config_parse_int(cfg, "quantizer", max_quantizer);
336         if (rret < 0) {
337             av_log(avctx, AV_LOG_ERROR, "Could not set max quantizer.\n");
338             ret = AVERROR_EXTERNAL;
339             goto end;
340         }
341
342         if (avctx->qmin >= 0) {
343             rret = rav1e_config_parse_int(cfg, "min_quantizer", avctx->qmin);
344             if (rret < 0) {
345                 av_log(avctx, AV_LOG_ERROR, "Could not set min quantizer.\n");
346                 ret = AVERROR_EXTERNAL;
347                 goto end;
348             }
349         }
350
351         rret = rav1e_config_parse_int(cfg, "bitrate", avctx->bit_rate);
352         if (rret < 0) {
353             av_log(avctx, AV_LOG_ERROR, "Could not set bitrate.\n");
354             ret = AVERROR_INVALIDDATA;
355             goto end;
356         }
357     } else if (ctx->quantizer >= 0) {
358         if (avctx->bit_rate)
359             av_log(avctx, AV_LOG_WARNING, "Both bitrate and quantizer specified. Using quantizer mode.");
360
361         rret = rav1e_config_parse_int(cfg, "quantizer", ctx->quantizer);
362         if (rret < 0) {
363             av_log(avctx, AV_LOG_ERROR, "Could not set quantizer.\n");
364             ret = AVERROR_EXTERNAL;
365             goto end;
366         }
367     }
368
369     rret = rav1e_config_set_pixel_format(cfg, desc->comp[0].depth,
370                                          pix_fmt_map(avctx->pix_fmt),
371                                          chroma_loc_map(avctx->chroma_sample_location),
372                                          range_map(avctx->pix_fmt, avctx->color_range));
373     if (rret < 0) {
374         av_log(avctx, AV_LOG_ERROR, "Failed to set pixel format properties.\n");
375         ret = AVERROR_INVALIDDATA;
376         goto end;
377     }
378
379     /* rav1e's colorspace enums match standard values. */
380     rret = rav1e_config_set_color_description(cfg, (RaMatrixCoefficients) avctx->colorspace,
381                                               (RaColorPrimaries) avctx->color_primaries,
382                                               (RaTransferCharacteristics) avctx->color_trc);
383     if (rret < 0) {
384         av_log(avctx, AV_LOG_WARNING, "Failed to set color properties.\n");
385         if (avctx->err_recognition & AV_EF_EXPLODE) {
386             ret = AVERROR_INVALIDDATA;
387             goto end;
388         }
389     }
390
391     ctx->ctx = rav1e_context_new(cfg);
392     if (!ctx->ctx) {
393         av_log(avctx, AV_LOG_ERROR, "Failed to create rav1e encode context.\n");
394         ret = AVERROR_EXTERNAL;
395         goto end;
396     }
397
398     ret = 0;
399
400 end:
401
402     rav1e_config_unref(cfg);
403
404     return ret;
405 }
406
407 static int librav1e_send_frame(AVCodecContext *avctx, const AVFrame *frame)
408 {
409     librav1eContext *ctx = avctx->priv_data;
410     RaFrame *rframe = NULL;
411     int ret;
412
413     if (frame) {
414         const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(frame->format);
415
416         rframe = rav1e_frame_new(ctx->ctx);
417         if (!rframe) {
418             av_log(avctx, AV_LOG_ERROR, "Could not allocate new rav1e frame.\n");
419             return AVERROR(ENOMEM);
420         }
421
422         for (int i = 0; i < desc->nb_components; i++) {
423             int shift = i ? desc->log2_chroma_h : 0;
424             int bytes = desc->comp[0].depth == 8 ? 1 : 2;
425             rav1e_frame_fill_plane(rframe, i, frame->data[i],
426                                    (frame->height >> shift) * frame->linesize[i],
427                                    frame->linesize[i], bytes);
428         }
429     }
430
431     ret = rav1e_send_frame(ctx->ctx, rframe);
432     if (rframe)
433          rav1e_frame_unref(rframe); /* No need to unref if flushing. */
434
435     switch (ret) {
436     case RA_ENCODER_STATUS_SUCCESS:
437         break;
438     case RA_ENCODER_STATUS_ENOUGH_DATA:
439         return AVERROR(EAGAIN);
440     case RA_ENCODER_STATUS_FAILURE:
441         av_log(avctx, AV_LOG_ERROR, "Could not send frame: %s\n", rav1e_status_to_str(ret));
442         return AVERROR_EXTERNAL;
443     default:
444         av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_send_frame: %s\n", ret, rav1e_status_to_str(ret));
445         return AVERROR_UNKNOWN;
446     }
447
448     return 0;
449 }
450
451 static int librav1e_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
452 {
453     librav1eContext *ctx = avctx->priv_data;
454     RaPacket *rpkt = NULL;
455     int ret;
456
457 retry:
458
459     if (avctx->flags & AV_CODEC_FLAG_PASS1) {
460         int sret = get_stats(avctx, 0);
461         if (sret < 0)
462             return sret;
463     } else if (avctx->flags & AV_CODEC_FLAG_PASS2) {
464         int sret = set_stats(avctx);
465         if (sret < 0)
466             return sret;
467     }
468
469     ret = rav1e_receive_packet(ctx->ctx, &rpkt);
470     switch (ret) {
471     case RA_ENCODER_STATUS_SUCCESS:
472         break;
473     case RA_ENCODER_STATUS_LIMIT_REACHED:
474         if (avctx->flags & AV_CODEC_FLAG_PASS1) {
475             int sret = get_stats(avctx, 1);
476             if (sret < 0)
477                 return sret;
478         }
479         return AVERROR_EOF;
480     case RA_ENCODER_STATUS_ENCODED:
481         if (avctx->internal->draining)
482             goto retry;
483         return AVERROR(EAGAIN);
484     case RA_ENCODER_STATUS_NEED_MORE_DATA:
485         if (avctx->internal->draining) {
486             av_log(avctx, AV_LOG_ERROR, "Unexpected error when receiving packet after EOF.\n");
487             return AVERROR_EXTERNAL;
488         }
489         return AVERROR(EAGAIN);
490     case RA_ENCODER_STATUS_FAILURE:
491         av_log(avctx, AV_LOG_ERROR, "Could not encode frame: %s\n", rav1e_status_to_str(ret));
492         return AVERROR_EXTERNAL;
493     default:
494         av_log(avctx, AV_LOG_ERROR, "Unknown return code %d from rav1e_receive_packet: %s\n", ret, rav1e_status_to_str(ret));
495         return AVERROR_UNKNOWN;
496     }
497
498     ret = av_new_packet(pkt, rpkt->len);
499     if (ret < 0) {
500         av_log(avctx, AV_LOG_ERROR, "Could not allocate packet.\n");
501         rav1e_packet_unref(rpkt);
502         return ret;
503     }
504
505     memcpy(pkt->data, rpkt->data, rpkt->len);
506
507     if (rpkt->frame_type == RA_FRAME_TYPE_KEY)
508         pkt->flags |= AV_PKT_FLAG_KEY;
509
510     pkt->pts = pkt->dts = rpkt->input_frameno * avctx->ticks_per_frame;
511
512     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
513         int ret = av_bsf_send_packet(ctx->bsf, pkt);
514         if (ret < 0) {
515             av_log(avctx, AV_LOG_ERROR, "extradata extraction send failed.\n");
516             rav1e_packet_unref(rpkt);
517             av_packet_unref(pkt);
518             return ret;
519         }
520
521         ret = av_bsf_receive_packet(ctx->bsf, pkt);
522         if (ret < 0) {
523             av_log(avctx, AV_LOG_ERROR, "extradata extraction receive failed.\n");
524             rav1e_packet_unref(rpkt);
525             av_packet_unref(pkt);
526             return ret;
527         }
528     }
529
530     return 0;
531 }
532
533 #define OFFSET(x) offsetof(librav1eContext, x)
534 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
535
536 static const AVOption options[] = {
537     { "qp", "use constant quantizer mode", OFFSET(quantizer), AV_OPT_TYPE_INT, { .i64 = 100 }, -1, 255, VE },
538     { "speed", "what speed preset to use", OFFSET(speed), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 10, VE },
539     { "tiles", "number of tiles encode with", OFFSET(tiles), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
540     { "tile-rows", "number of tiles rows to encode with", OFFSET(tile_rows), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
541     { "tile-columns", "number of tiles columns to encode with", OFFSET(tile_cols), AV_OPT_TYPE_INT, { .i64 = 0 }, -1, INT64_MAX, VE },
542     { "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 },
543     { NULL }
544 };
545
546 static const AVCodecDefault librav1e_defaults[] = {
547     { "b",           "0" },
548     { "g",           "0" },
549     { "keyint_min",  "0" },
550     { "qmax",       "-1" },
551     { "qmin",       "-1" },
552     { NULL }
553 };
554
555 const enum AVPixelFormat librav1e_pix_fmts[] = {
556     AV_PIX_FMT_YUV420P,
557     AV_PIX_FMT_YUVJ420P,
558     AV_PIX_FMT_YUV420P10,
559     AV_PIX_FMT_YUV420P12,
560     AV_PIX_FMT_YUV422P,
561     AV_PIX_FMT_YUVJ422P,
562     AV_PIX_FMT_YUV422P10,
563     AV_PIX_FMT_YUV422P12,
564     AV_PIX_FMT_YUV444P,
565     AV_PIX_FMT_YUVJ444P,
566     AV_PIX_FMT_YUV444P10,
567     AV_PIX_FMT_YUV444P12,
568     AV_PIX_FMT_NONE
569 };
570
571 static const AVClass class = {
572     .class_name = "librav1e",
573     .item_name  = av_default_item_name,
574     .option     = options,
575     .version    = LIBAVUTIL_VERSION_INT,
576 };
577
578 AVCodec ff_librav1e_encoder = {
579     .name           = "librav1e",
580     .long_name      = NULL_IF_CONFIG_SMALL("librav1e AV1"),
581     .type           = AVMEDIA_TYPE_VIDEO,
582     .id             = AV_CODEC_ID_AV1,
583     .init           = librav1e_encode_init,
584     .send_frame     = librav1e_send_frame,
585     .receive_packet = librav1e_receive_packet,
586     .close          = librav1e_encode_close,
587     .priv_data_size = sizeof(librav1eContext),
588     .priv_class     = &class,
589     .defaults       = librav1e_defaults,
590     .pix_fmts       = librav1e_pix_fmts,
591     .capabilities   = AV_CODEC_CAP_DELAY | AV_CODEC_CAP_AUTO_THREADS,
592     .caps_internal  = FF_CODEC_CAP_INIT_CLEANUP,
593     .wrapper_name   = "librav1e",
594 };