From 796669d4834b490095ebad3e898bb90502bd02d4 Mon Sep 17 00:00:00 2001
From: Niklas Haas <git@haasn.dev>
Date: Wed, 25 Oct 2023 13:42:24 +0200
Subject: [PATCH] avfilter/vf_format: add color_ranges/spaces

Needed for fftools/ffmpeg_filter to be able to force specific output
formats via the filter chain.
---
 doc/filters.texi        |  8 ++++++++
 libavfilter/vf_format.c | 42 ++++++++++++++++++++++++++++++++++++++---
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/doc/filters.texi b/doc/filters.texi
index 51bd20b137..a63a635ec8 100644
--- a/doc/filters.texi
+++ b/doc/filters.texi
@@ -14323,6 +14323,14 @@ It accepts the following parameters:
 A '|'-separated list of pixel format names, such as
 "pix_fmts=yuv420p|monow|rgb24".
 
+@item color_spaces
+A '|'-separated list of color space names, such as
+"color_spaces=bt709|bt470bg|bt2020nc".
+
+@item color_ranges
+A '|'-separated list of color range names, such as
+"color_spaces=tv|pc".
+
 @end table
 
 @subsection Examples
diff --git a/libavfilter/vf_format.c b/libavfilter/vf_format.c
index 59832b7768..d1bb9daa97 100644
--- a/libavfilter/vf_format.c
+++ b/libavfilter/vf_format.c
@@ -40,14 +40,20 @@
 typedef struct FormatContext {
     const AVClass *class;
     char *pix_fmts;
+    char *csps;
+    char *ranges;
 
     AVFilterFormats *formats; ///< parsed from `pix_fmts`
+    AVFilterFormats *color_spaces; ///< parsed from `csps`
+    AVFilterFormats *color_ranges; ///< parsed from `ranges`
 } FormatContext;
 
 static av_cold void uninit(AVFilterContext *ctx)
 {
     FormatContext *s = ctx->priv;
     ff_formats_unref(&s->formats);
+    ff_formats_unref(&s->color_spaces);
+    ff_formats_unref(&s->color_ranges);
 }
 
 static av_cold int invert_formats(AVFilterFormats **fmts,
@@ -99,13 +105,35 @@ static av_cold int init(AVFilterContext *ctx)
             return ret;
     }
 
+    for (char *sep, *cur = s->csps; cur; cur = sep) {
+        sep = strchr(cur, '|');
+        if (sep && *sep)
+            *sep++ = 0;
+        if ((ret = av_color_space_from_name(cur)) < 0 ||
+            (ret = ff_add_format(&s->color_spaces, ret)) < 0)
+            return ret;
+    }
+
+    for (char *sep, *cur = s->ranges; cur; cur = sep) {
+        sep = strchr(cur, '|');
+        if (sep && *sep)
+            *sep++ = 0;
+        if ((ret = av_color_range_from_name(cur)) < 0 ||
+            (ret = ff_add_format(&s->color_ranges, ret)) < 0)
+            return ret;
+    }
+
     if (!strcmp(ctx->filter->name, "noformat")) {
-        if ((ret = invert_formats(&s->formats, ff_all_formats(AVMEDIA_TYPE_VIDEO))) < 0)
+        if ((ret = invert_formats(&s->formats,      ff_all_formats(AVMEDIA_TYPE_VIDEO))) < 0 ||
+            (ret = invert_formats(&s->color_spaces, ff_all_color_spaces())) < 0 ||
+            (ret = invert_formats(&s->color_ranges, ff_all_color_ranges())) < 0)
             return ret;
     }
 
     /* hold on to a ref for the lifetime of the filter */
-    if ((ret = ff_formats_ref(s->formats, &s->formats)) < 0)
+    if ((ret = ff_formats_ref(s->formats, &s->formats)) < 0 ||
+        s->color_spaces && (ret = ff_formats_ref(s->color_spaces, &s->color_spaces)) < 0 ||
+        s->color_ranges && (ret = ff_formats_ref(s->color_ranges, &s->color_ranges)) < 0)
         return ret;
 
     return 0;
@@ -114,14 +142,22 @@ static av_cold int init(AVFilterContext *ctx)
 static int query_formats(AVFilterContext *ctx)
 {
     FormatContext *s = ctx->priv;
+    int ret;
+
+    if ((ret = ff_set_common_formats(ctx, s->formats)) < 0 ||
+        s->color_spaces && (ret = ff_set_common_color_spaces(ctx, s->color_spaces)) < 0 ||
+        s->color_ranges && (ret = ff_set_common_color_ranges(ctx, s->color_ranges)) < 0)
+        return ret;
 
-    return ff_set_common_formats(ctx, s->formats);
+    return 0;
 }
 
 
 #define OFFSET(x) offsetof(FormatContext, x)
 static const AVOption options[] = {
     { "pix_fmts", "A '|'-separated list of pixel formats", OFFSET(pix_fmts), AV_OPT_TYPE_STRING, .flags = AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_FILTERING_PARAM },
+    { "color_spaces", "A '|'-separated list of color spaces", OFFSET(csps), AV_OPT_TYPE_STRING, .flags = AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_FILTERING_PARAM },
+    { "color_ranges", "A '|'-separated list of color ranges", OFFSET(ranges), AV_OPT_TYPE_STRING, .flags = AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_FILTERING_PARAM },
     { NULL }
 };