From 64e3fc906971e18ec9e2810826ac9bb285ea08bb Mon Sep 17 00:00:00 2001
From: Timo Rothenpieler <timo@rothenpieler.org>
Date: Sun, 31 Mar 2024 20:10:45 +0200
Subject: [PATCH] avcodec/nvenc: add support for HEVC temporal filtering

---
 libavcodec/nvenc.c      | 27 +++++++++++++++++++++++++++
 libavcodec/nvenc.h      |  2 ++
 libavcodec/nvenc_hevc.c |  6 ++++++
 3 files changed, 35 insertions(+)

diff --git a/libavcodec/nvenc.c b/libavcodec/nvenc.c
index 8327496937..5afd7bf218 100644
--- a/libavcodec/nvenc.c
+++ b/libavcodec/nvenc.c
@@ -594,6 +594,14 @@ static int nvenc_check_capabilities(AVCodecContext *avctx)
         return AVERROR(ENOSYS);
     }
 
+#ifdef NVENC_HAVE_TEMPORAL_FILTER
+    ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_TEMPORAL_FILTER);
+    if(ctx->tf_level > 0 && ret <= 0) {
+        av_log(avctx, AV_LOG_WARNING, "Temporal filtering not supported by the device\n");
+        return AVERROR(ENOSYS);
+    }
+#endif
+
     ctx->support_dyn_bitrate = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_DYN_BITRATE_CHANGE);
 
     return 0;
@@ -1396,6 +1404,25 @@ static av_cold int nvenc_setup_hevc_config(AVCodecContext *avctx)
     hevc->numRefL1 = avctx->refs;
 #endif
 
+#ifdef NVENC_HAVE_TEMPORAL_FILTER
+    if (ctx->tf_level >= 0) {
+        hevc->tfLevel = ctx->tf_level;
+
+        switch (ctx->tf_level)
+        {
+            case NV_ENC_TEMPORAL_FILTER_LEVEL_0:
+            case NV_ENC_TEMPORAL_FILTER_LEVEL_4:
+                break;
+            default:
+                av_log(avctx, AV_LOG_ERROR, "Invalid temporal filtering level.\n");
+                return AVERROR(EINVAL);
+        }
+
+        if (ctx->encode_config.frameIntervalP < 5)
+            av_log(avctx, AV_LOG_WARNING, "Temporal filtering needs at least 4 B-Frames (-bf 4).\n");
+    }
+#endif
+
     return 0;
 }
 
diff --git a/libavcodec/nvenc.h b/libavcodec/nvenc.h
index d99d8a0d76..c320c2514f 100644
--- a/libavcodec/nvenc.h
+++ b/libavcodec/nvenc.h
@@ -86,6 +86,7 @@ typedef void ID3D11Device;
 // SDK 12.2 compile time feature checks
 #if NVENCAPI_CHECK_VERSION(12, 2)
 #define NVENC_HAVE_NEW_BIT_DEPTH_API
+#define NVENC_HAVE_TEMPORAL_FILTER
 #endif
 
 typedef struct NvencSurface
@@ -271,6 +272,7 @@ typedef struct NvencContext
     int highbitdepth;
     int max_slice_size;
     int rgb_mode;
+    int tf_level;
 } NvencContext;
 
 int ff_nvenc_encode_init(AVCodecContext *avctx);
diff --git a/libavcodec/nvenc_hevc.c b/libavcodec/nvenc_hevc.c
index 65fcb4efb8..8559aa6cfb 100644
--- a/libavcodec/nvenc_hevc.c
+++ b/libavcodec/nvenc_hevc.c
@@ -196,6 +196,12 @@ static const AVOption options[] = {
                                                             OFFSET(max_slice_size), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INT_MAX, VE },
     { "constrained-encoding", "Enable constrainedFrame encoding where each slice in the constrained picture is independent of other slices",
                                                             OFFSET(constrained_encoding), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
+#ifdef NVENC_HAVE_TEMPORAL_FILTER
+    { "tf_level",     "Specifies the strength of the temporal filtering",
+                                                            OFFSET(tf_level),     AV_OPT_TYPE_INT,   { .i64 = -1 }, -1, INT_MAX, VE, .unit = "tf_level" },
+    { "0",            "",                                   0,                    AV_OPT_TYPE_CONST, { .i64 = NV_ENC_TEMPORAL_FILTER_LEVEL_0 }, 0, 0, VE, .unit = "tf_level" },
+    { "4",            "",                                   0,                    AV_OPT_TYPE_CONST, { .i64 = NV_ENC_TEMPORAL_FILTER_LEVEL_4 }, 0, 0, VE, .unit = "tf_level" },
+#endif
     { NULL }
 };