[FFmpeg-devel] [PATCH] avcodec/libsvtav1: Add support for multipass encoding

[Gustav Grusell]

Implements support for 2-pass CRF and 3-pass VBR by implementing
reading and writing of stats file, and passing the pass number on
to the encoder. For 3-pass VBR, the first pass should be run with
'-pass 1', the second with '-pass 3' and the third with '-pass 2'.

Signed-off-by: Gustav Grusell <gustav.grusell at gmail.com>
---
 libavcodec/libsvtav1.c | 82 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/libavcodec/libsvtav1.c b/libavcodec/libsvtav1.c
index 2f5634cee0..73fece49b7 100644
--- a/libavcodec/libsvtav1.c
+++ b/libavcodec/libsvtav1.c
@@ -24,6 +24,7 @@
 #include <EbSvtAv1ErrorCodes.h>
 #include <EbSvtAv1Enc.h>
 
+#include "libavutil/base64.h"
 #include "libavutil/common.h"
 #include "libavutil/frame.h"
 #include "libavutil/imgutils.h"
@@ -312,6 +313,22 @@ static int config_enc_params(EbSvtAv1EncConfiguration *param,
         cpb_props->avg_bitrate = avctx->bit_rate;
     }
 
+    if (avctx->flags & AV_CODEC_FLAG_PASS2) {
+        if (param->rate_control_mode == SVT_AV1_RC_MODE_VBR) {
+            if (avctx->flags & AV_CODEC_FLAG_PASS1) {
+                param->pass = 2;
+            } else {
+                param->pass = 3;
+            }
+        } else {
+            param->pass = 2;
+        }
+    } else if (avctx->flags & AV_CODEC_FLAG_PASS1) {
+        param->pass = 1;
+    } else {
+        param->pass = 0;
+    }
+
     return 0;
 }
 
@@ -371,6 +388,34 @@ static av_cold int eb_enc_init(AVCodecContext *avctx)
         return ret;
     }
 
+    if (svt_enc->enc_params.pass >= 2) {
+        int decode_size;
+
+        if (!avctx->stats_in) {
+            av_log(avctx, AV_LOG_ERROR, "No stats file for %s pass\n",
+                   svt_enc->enc_params.pass == 2 ? "second" : "third");
+            return AVERROR_INVALIDDATA;
+        }
+
+        svt_enc->enc_params.rc_stats_buffer.sz = strlen(avctx->stats_in) * 3 / 4;
+        svt_enc->enc_params.rc_stats_buffer.buf = av_malloc(svt_enc->enc_params.rc_stats_buffer.sz);
+        if (!svt_enc->enc_params.rc_stats_buffer.buf) {
+            av_log(avctx, AV_LOG_ERROR,
+                   "Stat buffer alloc (%"SIZE_SPECIFIER" bytes) failed\n",
+                   svt_enc->enc_params.rc_stats_buffer.sz);
+            svt_enc->enc_params.rc_stats_buffer.sz = 0;
+            return AVERROR(ENOMEM);
+        }
+        decode_size = av_base64_decode(svt_enc->enc_params.rc_stats_buffer.buf, avctx->stats_in,
+                                       svt_enc->enc_params.rc_stats_buffer.sz);
+        if (decode_size < 0) {
+            av_log(avctx, AV_LOG_ERROR, "Stat buffer decode failed\n");
+            return AVERROR_INVALIDDATA;
+        }
+
+        svt_enc->enc_params.rc_stats_buffer.sz = decode_size;
+    }
+
     svt_ret = svt_av1_enc_set_parameter(svt_enc->svt_handle, &svt_enc->enc_params);
     if (svt_ret != EB_ErrorNone) {
         return svt_print_error(avctx, svt_ret, "Error setting encoder parameters");
@@ -544,6 +589,38 @@ static int eb_receive_packet(AVCodecContext *avctx, AVPacket *pkt)
     if (headerPtr->flags & EB_BUFFERFLAG_EOS)
         svt_enc->eos_flag = EOS_RECEIVED;
 
+    if (svt_enc->eos_flag == EOS_RECEIVED) {
+        if (svt_enc->enc_params.pass == 1) {
+            SvtAv1FixedBuf first_pass_stat;
+            EbErrorType    ret = svt_av1_enc_get_stream_info(
+                svt_enc->svt_handle,
+                SVT_AV1_STREAM_INFO_FIRST_PASS_STATS_OUT,
+                &first_pass_stat);
+            if (ret == EB_ErrorNone) {
+                size_t b64_size = AV_BASE64_SIZE(first_pass_stat.sz);
+                avctx->stats_out = av_malloc(b64_size);
+                if (!avctx->stats_out) {
+                    av_log(avctx, AV_LOG_ERROR, "Stat buffer alloc (%"SIZE_SPECIFIER" bytes) failed\n",
+                           b64_size);
+                    return AVERROR(ENOMEM);
+                }
+                av_base64_encode(avctx->stats_out, b64_size, first_pass_stat.buf,
+                                 first_pass_stat.sz);
+            }
+        }
+        if (svt_enc->enc_params.pass == 2) {
+            size_t b64_size = AV_BASE64_SIZE(svt_enc->enc_params.rc_stats_buffer.sz);
+            avctx->stats_out = av_malloc(b64_size);
+            if (!avctx->stats_out) {
+                av_log(avctx, AV_LOG_ERROR, "Stat buffer alloc (%"SIZE_SPECIFIER" bytes) failed\n",
+                       b64_size);
+                return AVERROR(ENOMEM);
+            }
+            av_base64_encode(avctx->stats_out, b64_size, svt_enc->enc_params.rc_stats_buffer.buf,
+                             svt_enc->enc_params.rc_stats_buffer.sz);
+        }
+    }
+
     ff_side_data_set_encoder_stats(pkt, headerPtr->qp * FF_QP2LAMBDA, NULL, 0, pict_type);
 
     svt_av1_enc_release_out_buffer(&headerPtr);
@@ -564,6 +641,11 @@ static av_cold int eb_enc_close(AVCodecContext *avctx)
         av_freep(&svt_enc->in_buf);
     }
 
+    if (svt_enc->enc_params.rc_stats_buffer.buf) {
+        av_freep(&svt_enc->enc_params.rc_stats_buffer.buf);
+        svt_enc->enc_params.rc_stats_buffer.sz = 0;
+    }
+
     av_buffer_pool_uninit(&svt_enc->pool);
     av_frame_free(&svt_enc->frame);
 
-- 
2.34.1