git.sesse.net Git - ffmpeg/blob - libavfilter/vf_fftfilt.c

   1 /*
   2  * Copyright (c) 2015 Arwa Arif <arwaarif1994@gmail.com>
   3  *
   4  * This file is part of FFmpeg.
   5  *
   6  * FFmpeg is free software; you can redistribute it and/or modify it
   7  * under the terms of the GNU Lesser General Public License as published
   8  * by the Free Software Foundation; either version 2.1 of the License,
   9  * or (at your option) any later version.
  10  *
  11  * FFmpeg is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with FFmpeg; if not, write to the Free Software
  18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  19  */
  20
  21 /**
  22  * @file
  23  * FFT domain filtering.
  24  */
  25
  26 #include "libavfilter/internal.h"
  27 #include "libavutil/common.h"
  28 #include "libavutil/imgutils.h"
  29 #include "libavutil/opt.h"
  30 #include "libavutil/pixdesc.h"
  31 #include "libavcodec/avfft.h"
  32 #include "libavutil/eval.h"
  33
  34 #define MAX_PLANES 4
  35
  36 typedef struct FFTFILTContext {
  37     const AVClass *class;
  38
  39     RDFTContext *rdft;
  40     int rdft_hbits[MAX_PLANES];
  41     int rdft_vbits[MAX_PLANES];
  42     size_t rdft_hlen[MAX_PLANES];
  43     size_t rdft_vlen[MAX_PLANES];
  44     FFTSample *rdft_hdata[MAX_PLANES];
  45     FFTSample *rdft_vdata[MAX_PLANES];
  46
  47     int dc[MAX_PLANES];
  48     char *weight_str[MAX_PLANES];
  49     AVExpr *weight_expr[MAX_PLANES];
  50     double *weight[MAX_PLANES];
  51
  52 } FFTFILTContext;
  53
  54 static const char *const var_names[] = {   "X",   "Y",   "W",   "H",     NULL    };
  55 enum                                   { VAR_X, VAR_Y, VAR_W, VAR_H, VAR_VARS_NB };
  56
  57 enum { Y = 0, U, V };
  58
  59 #define OFFSET(x) offsetof(FFTFILTContext, x)
  60 #define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
  61
  62 static const AVOption fftfilt_options[] = {
  63     { "dc_Y",  "adjust gain in Y plane",              OFFSET(dc[Y]),      AV_OPT_TYPE_INT,    {.i64 = 0},      0,     1000,     FLAGS },
  64     { "dc_U",  "adjust gain in U plane",              OFFSET(dc[U]),      AV_OPT_TYPE_INT,    {.i64 = 0},      0,     1000,     FLAGS },
  65     { "dc_V",  "adjust gain in V plane",              OFFSET(dc[V]),      AV_OPT_TYPE_INT,    {.i64 = 0},      0,     1000,     FLAGS },
  66     { "weight_Y", "set luminance expression in Y plane",   OFFSET(weight_str[Y]), AV_OPT_TYPE_STRING, {.str = "1"}, CHAR_MIN, CHAR_MAX, FLAGS },
  67     { "weight_U", "set chrominance expression in U plane", OFFSET(weight_str[U]), AV_OPT_TYPE_STRING, {.str = NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
  68     { "weight_V", "set chrominance expression in V plane", OFFSET(weight_str[V]), AV_OPT_TYPE_STRING, {.str = NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
  69     {NULL},
  70 };
  71
  72 AVFILTER_DEFINE_CLASS(fftfilt);
  73
  74 static inline double lum(void *priv, double x, double y, int plane)
  75 {
  76     FFTFILTContext *s = priv;
  77     return s->rdft_vdata[plane][(int)x * s->rdft_vlen[plane] + (int)y];
  78 }
  79
  80 static double weight_Y(void *priv, double x, double y) { return lum(priv, x, y, Y); }
  81 static double weight_U(void *priv, double x, double y) { return lum(priv, x, y, U); }
  82 static double weight_V(void *priv, double x, double y) { return lum(priv, x, y, V); }
  83
  84 static void copy_rev (FFTSample *dest, int w, int w2)
  85 {
  86     int i;
  87
  88     for (i = w; i < w + (w2-w)/2; i++)
  89         dest[i] = dest[2*w - i - 1];
  90
  91     for (; i < w2; i++)
  92         dest[i] = dest[w2 - i];
  93 }
  94
  95 /*Horizontal pass - RDFT*/
  96 static void rdft_horizontal(FFTFILTContext *s, AVFrame *in, int w, int h, int plane)
  97 {
  98     int i, j;
  99     s->rdft = av_rdft_init(s->rdft_hbits[plane], DFT_R2C);
 100
 101     for (i = 0; i < h; i++) {
 102         for (j = 0; j < w; j++)
 103             s->rdft_hdata[plane][i * s->rdft_hlen[plane] + j] = *(in->data[plane] + in->linesize[plane] * i + j);
 104
 105         copy_rev(s->rdft_hdata[plane] + i * s->rdft_hlen[plane], w, s->rdft_hlen[plane]);
 106     }
 107
 108     for (i = 0; i < h; i++)
 109         av_rdft_calc(s->rdft, s->rdft_hdata[plane] + i * s->rdft_hlen[plane]);
 110
 111     av_rdft_end(s->rdft);
 112 }
 113
 114 /*Vertical pass - RDFT*/
 115 static void rdft_vertical(FFTFILTContext *s, int h, int plane)
 116 {
 117     int i, j;
 118     s->rdft = av_rdft_init(s->rdft_vbits[plane], DFT_R2C);
 119
 120     for (i = 0; i < s->rdft_hlen[plane]; i++) {
 121         for (j = 0; j < h; j++)
 122             s->rdft_vdata[plane][i * s->rdft_vlen[plane] + j] =
 123             s->rdft_hdata[plane][j * s->rdft_hlen[plane] + i];
 124         copy_rev(s->rdft_vdata[plane] + i * s->rdft_vlen[plane], h, s->rdft_vlen[plane]);
 125     }
 126
 127     for (i = 0; i < s->rdft_hlen[plane]; i++)
 128         av_rdft_calc(s->rdft, s->rdft_vdata[plane] + i * s->rdft_vlen[plane]);
 129
 130     av_rdft_end(s->rdft);
 131 }
 132 /*Vertical pass - IRDFT*/
 133 static void irdft_vertical(FFTFILTContext *s, int h, int plane)
 134 {
 135     int i, j;
 136     s->rdft = av_rdft_init(s->rdft_vbits[plane], IDFT_C2R);
 137     for (i = 0; i < s->rdft_hlen[plane]; i++)
 138         av_rdft_calc(s->rdft, s->rdft_vdata[plane] + i * s->rdft_vlen[plane]);
 139
 140     for (i = 0; i < s->rdft_hlen[plane]; i++)
 141         for (j = 0; j < h; j++)
 142             s->rdft_hdata[plane][j * s->rdft_hlen[plane] + i] =
 143             s->rdft_vdata[plane][i * s->rdft_vlen[plane] + j];
 144
 145     av_rdft_end(s->rdft);
 146 }
 147
 148 /*Horizontal pass - IRDFT*/
 149 static void irdft_horizontal(FFTFILTContext *s, AVFrame *out, int w, int h, int plane)
 150 {
 151     int i, j;
 152     s->rdft = av_rdft_init(s->rdft_hbits[plane], IDFT_C2R);
 153     for (i = 0; i < h; i++)
 154         av_rdft_calc(s->rdft, s->rdft_hdata[plane] + i * s->rdft_hlen[plane]);
 155
 156     for (i = 0; i < h; i++)
 157         for (j = 0; j < w; j++)
 158             *(out->data[plane] + out->linesize[plane] * i + j) = av_clip(s->rdft_hdata[plane][i
 159                                                                          *s->rdft_hlen[plane] + j] * 4 /
 160                                                                          (s->rdft_hlen[plane] *
 161                                                                           s->rdft_vlen[plane]), 0, 255);
 162
 163     av_rdft_end(s->rdft);
 164 }
 165
 166 static av_cold int initialize(AVFilterContext *ctx)
 167 {
 168     FFTFILTContext *s = ctx->priv;
 169     int ret = 0, plane;
 170
 171     if (!s->dc[U] && !s->dc[V]) {
 172         s->dc[U] = s->dc[Y];
 173         s->dc[V] = s->dc[Y];
 174     } else {
 175         if (!s->dc[U]) s->dc[U] = s->dc[V];
 176         if (!s->dc[V]) s->dc[V] = s->dc[U];
 177     }
 178
 179     if (!s->weight_str[U] && !s->weight_str[V]) {
 180         s->weight_str[U] = av_strdup(s->weight_str[Y]);
 181         s->weight_str[V] = av_strdup(s->weight_str[Y]);
 182     } else {
 183         if (!s->weight_str[U]) s->weight_str[U] = av_strdup(s->weight_str[V]);
 184         if (!s->weight_str[V]) s->weight_str[V] = av_strdup(s->weight_str[U]);
 185     }
 186
 187     for (plane = 0; plane < 3; plane++) {
 188         static double (*p[])(void *, double, double) = { weight_Y, weight_U, weight_V };
 189         const char *const func2_names[] = {"weight_Y", "weight_U", "weight_V", NULL };
 190         double (*func2[])(void *, double, double) = { weight_Y, weight_U, weight_V, p[plane], NULL };
 191
 192         ret = av_expr_parse(&s->weight_expr[plane], s->weight_str[plane], var_names,
 193                             NULL, NULL, func2_names, func2, 0, ctx);
 194         if (ret < 0)
 195             break;
 196     }
 197     return ret;
 198 }
 199
 200 static int config_props(AVFilterLink *inlink)
 201 {
 202     FFTFILTContext *s = inlink->dst->priv;
 203     const AVPixFmtDescriptor *desc;
 204     int rdft_hbits, rdft_vbits, i, j, plane;
 205     double values[VAR_VARS_NB];
 206
 207     desc = av_pix_fmt_desc_get(inlink->format);
 208     for (i = 0; i < desc->nb_components; i++) {
 209         int w = inlink->w;
 210         int h = inlink->h;
 211
 212         /* RDFT - Array initialization for Horizontal pass*/
 213         for (rdft_hbits = 1; 1 << rdft_hbits < w*10/9; rdft_hbits++);
 214         s->rdft_hbits[i] = rdft_hbits;
 215         s->rdft_hlen[i] = 1 << rdft_hbits;
 216         if (!(s->rdft_hdata[i] = av_malloc_array(h, s->rdft_hlen[i] * sizeof(FFTSample))))
 217             return AVERROR(ENOMEM);
 218
 219         /* RDFT - Array initialization for Vertical pass*/
 220         for (rdft_vbits = 1; 1 << rdft_vbits < h*10/9; rdft_vbits++);
 221         s->rdft_vbits[i] = rdft_vbits;
 222         s->rdft_vlen[i] = 1 << rdft_vbits;
 223         if (!(s->rdft_vdata[i] = av_malloc_array(s->rdft_hlen[i], s->rdft_vlen[i] * sizeof(FFTSample))))
 224             return AVERROR(ENOMEM);
 225     }
 226
 227     /*Luminance value - Array initialization*/
 228     values[VAR_W] = inlink->w;
 229     values[VAR_H] = inlink->h;
 230     for (plane = 0; plane < 3; plane++)
 231     {
 232         if(!(s->weight[plane] = av_malloc_array(s->rdft_hlen[plane], s->rdft_vlen[plane] * sizeof(double))))
 233             return AVERROR(ENOMEM);
 234         for (i = 0; i < s->rdft_hlen[plane]; i++)
 235         {
 236             values[VAR_X] = i;
 237             for (j = 0; j < s->rdft_vlen[plane]; j++)
 238             {
 239                 values[VAR_Y] = j;
 240                 s->weight[plane][i * s->rdft_vlen[plane] + j] =
 241                 av_expr_eval(s->weight_expr[plane], values, s);
 242             }
 243         }
 244     }
 245     return 0;
 246 }
 247
 248 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 249 {
 250     AVFilterContext *ctx = inlink->dst;
 251     AVFilterLink *outlink = inlink->dst->outputs[0];
 252     const AVPixFmtDescriptor *desc;
 253     FFTFILTContext *s = ctx->priv;
 254     AVFrame *out;
 255     int i, j, plane;
 256
 257     out = ff_get_video_buffer(outlink, inlink->w, inlink->h);
 258     if (!out)
 259         return AVERROR(ENOMEM);
 260
 261     av_frame_copy_props(out, in);
 262
 263     desc = av_pix_fmt_desc_get(inlink->format);
 264     for (plane = 0; plane < desc->nb_components; plane++) {
 265         int w = inlink->w;
 266         int h = inlink->h;
 267
 268         if (plane == 1 || plane == 2) {
 269             w = AV_CEIL_RSHIFT(w, desc->log2_chroma_w);
 270             h = AV_CEIL_RSHIFT(h, desc->log2_chroma_h);
 271         }
 272
 273         rdft_horizontal(s, in, w, h, plane);
 274         rdft_vertical(s, h, plane);
 275
 276         /*Change user defined parameters*/
 277         for (i = 0; i < s->rdft_hlen[plane]; i++)
 278             for (j = 0; j < s->rdft_vlen[plane]; j++)
 279                 s->rdft_vdata[plane][i * s->rdft_vlen[plane] + j] *=
 280                   s->weight[plane][i * s->rdft_vlen[plane] + j];
 281
 282         s->rdft_vdata[plane][0] += s->rdft_hlen[plane] * s->rdft_vlen[plane] * s->dc[plane];
 283
 284         irdft_vertical(s, h, plane);
 285         irdft_horizontal(s, out, w, h, plane);
 286     }
 287
 288     av_frame_free(&in);
 289     return ff_filter_frame(outlink, out);
 290 }
 291
 292 static av_cold void uninit(AVFilterContext *ctx)
 293 {
 294     FFTFILTContext *s = ctx->priv;
 295     int i;
 296     for (i = 0; i < MAX_PLANES; i++) {
 297         av_free(s->rdft_hdata[i]);
 298         av_free(s->rdft_vdata[i]);
 299         av_expr_free(s->weight_expr[i]);
 300         av_free(s->weight[i]);
 301     }
 302 }
 303
 304 static int query_formats(AVFilterContext *ctx)
 305 {
 306     static const enum AVPixelFormat pixel_fmts_fftfilt[] = {
 307         AV_PIX_FMT_GRAY8,
 308         AV_PIX_FMT_YUV444P,
 309         AV_PIX_FMT_NONE
 310     };
 311
 312     AVFilterFormats *fmts_list = ff_make_format_list(pixel_fmts_fftfilt);
 313     if (!fmts_list)
 314         return AVERROR(ENOMEM);
 315     return ff_set_common_formats(ctx, fmts_list);
 316 }
 317
 318 static const AVFilterPad fftfilt_inputs[] = {
 319     {
 320         .name = "default",
 321         .type = AVMEDIA_TYPE_VIDEO,
 322         .config_props = config_props,
 323         .filter_frame = filter_frame,
 324     },
 325     { NULL }
 326 };
 327
 328 static const AVFilterPad fftfilt_outputs[] = {
 329     {
 330         .name = "default",
 331         .type = AVMEDIA_TYPE_VIDEO,
 332     },
 333     { NULL }
 334 };
 335
 336 AVFilter ff_vf_fftfilt = {
 337     .name            = "fftfilt",
 338     .description     = NULL_IF_CONFIG_SMALL("Apply arbitrary expressions to pixels in frequency domain."),
 339     .priv_size       = sizeof(FFTFILTContext),
 340     .priv_class      = &fftfilt_class,
 341     .inputs          = fftfilt_inputs,
 342     .outputs         = fftfilt_outputs,
 343     .query_formats   = query_formats,
 344     .init            = initialize,
 345     .uninit          = uninit,
 346 };