lavfi/geq: improve support for formats with alpha plane

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

1 files changed, 28 insertions, 17 deletions

diff --git a/libavfilter/vf_geq.c b/libavfilter/vf_geq.c
index 2f7b50b647..72ecda0af1 100644
--- a/libavfilter/vf_geq.c
+++ b/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;