lavfi/geq: improve support for formats with alpha plane

Signed-off-by: Paul B Mahol <onemda@gmail.com>

doc/filters.texi

@@ -3047,7 +3047,7 @@ For more information see:
 
 @section geq
 
-The filter takes one, two or three equations as parameter, separated by ':'.
+The filter takes one, two, three or four equations as parameter, separated by ':'.
 The first equation is mandatory and applies to the luma plane. The two
 following are respectively for chroma blue and chroma red planes.
 
@@ -3060,11 +3060,14 @@ the luminance expression
 the chrominance blue expression
 @item cr_expr
 the chrominance red expression
+@item alpha_expr
+the alpha expression
 @end table
 
 If one of the chrominance expression is not defined, it falls back on the other
-one. If none of them are specified, they will evaluate the luminance
-expression.
+one. If no alpha expression is specified it will evaluate to opaque value.
+If none of chrominance expressions are
+specified, they will evaluate the luminance expression.
 
 The expressions can use the following variables and functions:
 
@@ -3097,11 +3100,15 @@ plane.
 
 @item cb(x, y)
 Return the value of the pixel at location (@var{x},@var{y}) of the
-blue-difference chroma plane.
+blue-difference chroma plane. Returns 0 if there is no such plane.
 
 @item cr(x, y)
 Return the value of the pixel at location (@var{x},@var{y}) of the
-red-difference chroma plane.
+red-difference chroma plane. Returns 0 if there is no such plane.
+
+@item alpha(x, y)
+Return the value of the pixel at location (@var{x},@var{y}) of the alpha
+plane. Returns 0 if there is no such plane.
 @end table
 
 For functions, if @var{x} and @var{y} are outside the area, the value will be

libavfilter/vf_geq.c

@@ -34,20 +34,22 @@
 
 typedef struct {
     const AVClass *class;
-    AVExpr *e[3];               ///< expressions for each plane
-    char *expr_str[3];          ///< expression strings for each plane
+    AVExpr *e[4];               ///< expressions for each plane
+    char *expr_str[4];          ///< expression strings for each plane
     int framenum;               ///< frame counter
     AVFilterBufferRef *picref;  ///< current input buffer
     int hsub, vsub;             ///< chroma subsampling
+    int planes;                 ///< number of planes
 } GEQContext;
 
 #define OFFSET(x) offsetof(GEQContext, x)
 #define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
 
 static const AVOption geq_options[] = {
-    { "lum_expr",   "set luminance expression",   OFFSET(expr_str),                   AV_OPT_TYPE_STRING, {.str=NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
-    { "cb_expr",    "set chroma blue expression", OFFSET(expr_str) +   sizeof(char*), AV_OPT_TYPE_STRING, {.str=NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
-    { "cr_expr",    "set chroma red expression",  OFFSET(expr_str) + 2*sizeof(char*), AV_OPT_TYPE_STRING, {.str=NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
+    { "lum_expr",   "set luminance expression",   OFFSET(expr_str[0]), AV_OPT_TYPE_STRING, {.str=NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
+    { "cb_expr",    "set chroma blue expression", OFFSET(expr_str[1]), AV_OPT_TYPE_STRING, {.str=NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
+    { "cr_expr",    "set chroma red expression",  OFFSET(expr_str[2]), AV_OPT_TYPE_STRING, {.str=NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
+    { "alpha_expr", "set alpha expression",       OFFSET(expr_str[3]), AV_OPT_TYPE_STRING, {.str=NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
     {NULL},
 };
 
@@ -60,8 +62,11 @@ static inline double getpix(void *priv, double x, double y, int plane)
     AVFilterBufferRef *picref = geq->picref;
     const uint8_t *src = picref->data[plane];
     const int linesize = picref->linesize[plane];
-    const int w = picref->video->w >> (plane ? geq->hsub : 0);
-    const int h = picref->video->h >> (plane ? geq->vsub : 0);
+    const int w = picref->video->w >> ((plane == 1 || plane == 2) ? geq->hsub : 0);
+    const int h = picref->video->h >> ((plane == 1 || plane == 2) ? geq->vsub : 0);
+
+    if (!src)
+        return 0;
 
     xi = x = av_clipf(x, 0, w - 2);
     yi = y = av_clipf(y, 0, h - 2);
@@ -78,6 +83,7 @@ static inline double getpix(void *priv, double x, double y, int plane)
 static double lum(void *priv, double x, double y) { return getpix(priv, x, y, 0); }
 static double  cb(void *priv, double x, double y) { return getpix(priv, x, y, 1); }
 static double  cr(void *priv, double x, double y) { return getpix(priv, x, y, 2); }
+static double alpha(void *priv, double x, double y) { return getpix(priv, x, y, 3); }
 
 static const char *const var_names[] = {   "X",   "Y",   "W",   "H",   "N",   "SW",   "SH",   "T",        NULL };
 enum                                   { VAR_X, VAR_Y, VAR_W, VAR_H, VAR_N, VAR_SW, VAR_SH, VAR_T, VAR_VARS_NB };
@@ -86,7 +92,7 @@ static av_cold int geq_init(AVFilterContext *ctx, const char *args)
 {
     GEQContext *geq = ctx->priv;
     int plane, ret = 0;
-    static const char *shorthand[] = { "lum_expr", "cb_expr", "cr_expr", NULL };
+    static const char *shorthand[] = { "lum_expr", "cb_expr", "cr_expr", "alpha_expr", NULL };
 
     geq->class = &geq_class;
     av_opt_set_defaults(geq);
@@ -110,15 +116,18 @@ static av_cold int geq_init(AVFilterContext *ctx, const char *args)
         if (!geq->expr_str[2]) geq->expr_str[2] = av_strdup(geq->expr_str[1]);
     }
 
-    if (!geq->expr_str[1] || !geq->expr_str[2]) {
+    if (!geq->expr_str[3])
+        geq->expr_str[3] = av_strdup("255");
+
+    if (!geq->expr_str[1] || !geq->expr_str[2] || !geq->expr_str[3]) {
         ret = AVERROR(ENOMEM);
         goto end;
     }
 
-    for (plane = 0; plane < 3; plane++) {
-        static double (*p[])(void *, double, double) = { lum, cb, cr };
-        static const char *const func2_names[]    = { "lum", "cb", "cr", "p", NULL };
-        double (*func2[])(void *, double, double) = { lum, cb, cr, p[plane], NULL };
+    for (plane = 0; plane < 4; plane++) {
+        static double (*p[])(void *, double, double) = { lum, cb, cr, alpha };
+        static const char *const func2_names[]    = { "lum", "cb", "cr", "alpha", "p", NULL };
+        double (*func2[])(void *, double, double) = { lum, cb, cr, alpha, p[plane], NULL };
 
         ret = av_expr_parse(&geq->e[plane], geq->expr_str[plane], var_names,
                             NULL, NULL, func2_names, func2, 0, ctx);
@@ -135,7 +144,8 @@ static int geq_query_formats(AVFilterContext *ctx)
     static const enum PixelFormat pix_fmts[] = {
         AV_PIX_FMT_YUV444P,  AV_PIX_FMT_YUV422P,  AV_PIX_FMT_YUV420P,
         AV_PIX_FMT_YUV411P,  AV_PIX_FMT_YUV410P,  AV_PIX_FMT_YUV440P,
-        AV_PIX_FMT_YUVA420P,
+        AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUVA420P,
+        AV_PIX_FMT_GRAY8,
         AV_PIX_FMT_NONE
     };
     ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
@@ -149,6 +159,7 @@ static int geq_config_props(AVFilterLink *inlink)
 
     geq->hsub = desc->log2_chroma_w;
     geq->vsub = desc->log2_chroma_h;
+    geq->planes = desc->nb_components;
     return 0;
 }
 
@@ -171,12 +182,12 @@ static int geq_filter_frame(AVFilterLink *inlink, AVFilterBufferRef *in)
     }
     avfilter_copy_buffer_ref_props(out, in);
 
-    for (plane = 0; plane < 3; plane++) {
+    for (plane = 0; plane < geq->planes && out->data[plane]; plane++) {
         int x, y;
         uint8_t *dst = out->data[plane];
         const int linesize = out->linesize[plane];
-        const int w = inlink->w >> (plane ? geq->hsub : 0);
-        const int h = inlink->h >> (plane ? geq->vsub : 0);
+        const int w = inlink->w >> ((plane == 1 || plane == 2) ? geq->hsub : 0);
+        const int h = inlink->h >> ((plane == 1 || plane == 2) ? geq->vsub : 0);
 
         values[VAR_W]  = w;
         values[VAR_H]  = h;