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 typedef struct {
  35     const AVClass *class;
  36
  37     RDFTContext *rdft;
  38     int rdft_hbits[3];
  39     int rdft_vbits[3];
  40     size_t rdft_hlen[3];
  41     size_t rdft_vlen[3];
  42     FFTSample *rdft_hdata[3];
  43     FFTSample *rdft_vdata[3];
  44
  45     int dc[3];
  46     char *weight_str[3];
  47     AVExpr *weight_expr[3];
  48     double *weight[3];
  49
  50 } FFTFILTContext;
  51
  52 static const char *const var_names[] = {   "X",   "Y",   "W",   "H",     NULL    };
  53 enum                                   { VAR_X, VAR_Y, VAR_W, VAR_H, VAR_VARS_NB };
  54
  55 enum { Y = 0, U, V };
  56
  57 #define OFFSET(x) offsetof(FFTFILTContext, x)
  58 #define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM
  59
  60 static const AVOption fftfilt_options[] = {
  61     { "dc_Y",  "adjust gain in Y plane",              OFFSET(dc[Y]),      AV_OPT_TYPE_INT,    {.i64 = 0},      0,     1000,     FLAGS },
  62     { "dc_U",  "adjust gain in U plane",              OFFSET(dc[U]),      AV_OPT_TYPE_INT,    {.i64 = 0},      0,     1000,     FLAGS },
  63     { "dc_V",  "adjust gain in V plane",              OFFSET(dc[V]),      AV_OPT_TYPE_INT,    {.i64 = 0},      0,     1000,     FLAGS },
  64     { "weight_Y", "set luminance expression in Y plane",   OFFSET(weight_str[Y]), AV_OPT_TYPE_STRING, {.str = "1"}, CHAR_MIN, CHAR_MAX, FLAGS },
  65     { "weight_U", "set chrominance expression in U plane", OFFSET(weight_str[U]), AV_OPT_TYPE_STRING, {.str = NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
  66     { "weight_V", "set chrominance expression in V plane", OFFSET(weight_str[V]), AV_OPT_TYPE_STRING, {.str = NULL}, CHAR_MIN, CHAR_MAX, FLAGS },
  67     {NULL},
  68 };
  69
  70 AVFILTER_DEFINE_CLASS(fftfilt);
  71
  72 static inline double lum(void *priv, double x, double y, int plane)
  73 {
  74     FFTFILTContext *fftfilt = priv;
  75     return fftfilt->rdft_vdata[plane][(int)x * fftfilt->rdft_vlen[plane] + (int)y];
  76 }
  77
  78 static double weight_Y(void *priv, double x, double y) { return lum(priv, x, y, Y); }
  79 static double weight_U(void *priv, double x, double y) { return lum(priv, x, y, U); }
  80 static double weight_V(void *priv, double x, double y) { return lum(priv, x, y, V); }
  81
  82 static void copy_rev (FFTSample *dest, int w, int w2)
  83 {
  84     int i;
  85
  86     for (i = w; i < w + (w2-w)/2; i++)
  87         dest[i] = dest[2*w - i - 1];
  88
  89     for (; i < w2; i++)
  90         dest[i] = dest[w2 - i];
  91 }
  92
  93 /*Horizontal pass - RDFT*/
  94 static void rdft_horizontal(FFTFILTContext *fftfilt, AVFrame *in, int w, int h, int plane)
  95 {
  96     int i, j;
  97     fftfilt->rdft = av_rdft_init(fftfilt->rdft_hbits[plane], DFT_R2C);
  98
  99     for (i = 0; i < h; i++) {
 100         for (j = 0; j < w; j++)
 101             fftfilt->rdft_hdata[plane][i * fftfilt->rdft_hlen[plane] + j] = *(in->data[plane] + in->linesize[plane] * i + j);
 102
 103         copy_rev(fftfilt->rdft_hdata[plane] + i * fftfilt->rdft_hlen[plane], w, fftfilt->rdft_hlen[plane]);
 104     }
 105
 106     for (i = 0; i < h; i++)
 107         av_rdft_calc(fftfilt->rdft, fftfilt->rdft_hdata[plane] + i * fftfilt->rdft_hlen[plane]);
 108
 109     av_rdft_end(fftfilt->rdft);
 110 }
 111
 112 /*Vertical pass - RDFT*/
 113 static void rdft_vertical(FFTFILTContext *fftfilt, int h, int plane)
 114 {
 115     int i, j;
 116     fftfilt->rdft = av_rdft_init(fftfilt->rdft_vbits[plane], DFT_R2C);
 117
 118     for (i = 0; i < fftfilt->rdft_hlen[plane]; i++) {
 119         for (j = 0; j < h; j++)
 120             fftfilt->rdft_vdata[plane][i * fftfilt->rdft_vlen[plane] + j] =
 121             fftfilt->rdft_hdata[plane][j * fftfilt->rdft_hlen[plane] + i];
 122         copy_rev(fftfilt->rdft_vdata[plane] + i * fftfilt->rdft_vlen[plane], h, fftfilt->rdft_vlen[plane]);
 123     }
 124
 125     for (i = 0; i < fftfilt->rdft_hlen[plane]; i++)
 126         av_rdft_calc(fftfilt->rdft, fftfilt->rdft_vdata[plane] + i * fftfilt->rdft_vlen[plane]);
 127
 128     av_rdft_end(fftfilt->rdft);
 129 }
 130 /*Vertical pass - IRDFT*/
 131 static void irdft_vertical(FFTFILTContext *fftfilt, int h, int plane)
 132 {
 133     int i, j;
 134     fftfilt->rdft = av_rdft_init(fftfilt->rdft_vbits[plane], IDFT_C2R);
 135     for (i = 0; i < fftfilt->rdft_hlen[plane]; i++)
 136         av_rdft_calc(fftfilt->rdft, fftfilt->rdft_vdata[plane] + i * fftfilt->rdft_vlen[plane]);
 137
 138     for (i = 0; i < fftfilt->rdft_hlen[plane]; i++)
 139         for (j = 0; j < h; j++)
 140             fftfilt->rdft_hdata[plane][j * fftfilt->rdft_hlen[plane] + i] =
 141             fftfilt->rdft_vdata[plane][i * fftfilt->rdft_vlen[plane] + j];
 142
 143     av_rdft_end(fftfilt->rdft);
 144 }
 145
 146 /*Horizontal pass - IRDFT*/
 147 static void irdft_horizontal(FFTFILTContext *fftfilt, AVFrame *out, int w, int h, int plane)
 148 {
 149     int i, j;
 150     fftfilt->rdft = av_rdft_init(fftfilt->rdft_hbits[plane], IDFT_C2R);
 151     for (i = 0; i < h; i++)
 152         av_rdft_calc(fftfilt->rdft, fftfilt->rdft_hdata[plane] + i * fftfilt->rdft_hlen[plane]);
 153
 154     for (i = 0; i < h; i++)
 155         for (j = 0; j < w; j++)
 156             *(out->data[plane] + out->linesize[plane] * i + j) = av_clip(fftfilt->rdft_hdata[plane][i
 157                                                                          *fftfilt->rdft_hlen[plane] + j] * 4 /
 158                                                                          (fftfilt->rdft_hlen[plane] *
 159                                                                           fftfilt->rdft_vlen[plane]), 0, 255);
 160
 161     av_rdft_end(fftfilt->rdft);
 162 }
 163
 164 static av_cold int initialize(AVFilterContext *ctx)
 165 {
 166     FFTFILTContext *fftfilt = ctx->priv;
 167     int ret = 0, plane;
 168
 169     if (!fftfilt->dc[U] && !fftfilt->dc[V]) {
 170         fftfilt->dc[U] = fftfilt->dc[Y];
 171         fftfilt->dc[V] = fftfilt->dc[Y];
 172     } else {
 173         if (!fftfilt->dc[U]) fftfilt->dc[U] = fftfilt->dc[V];
 174         if (!fftfilt->dc[V]) fftfilt->dc[V] = fftfilt->dc[U];
 175     }
 176
 177     if (!fftfilt->weight_str[U] && !fftfilt->weight_str[V]) {
 178         fftfilt->weight_str[U] = av_strdup(fftfilt->weight_str[Y]);
 179         fftfilt->weight_str[V] = av_strdup(fftfilt->weight_str[Y]);
 180     } else {
 181         if (!fftfilt->weight_str[U]) fftfilt->weight_str[U] = av_strdup(fftfilt->weight_str[V]);
 182         if (!fftfilt->weight_str[V]) fftfilt->weight_str[V] = av_strdup(fftfilt->weight_str[U]);
 183     }
 184
 185     for (plane = 0; plane < 3; plane++) {
 186         static double (*p[])(void *, double, double) = { weight_Y, weight_U, weight_V };
 187         const char *const func2_names[] = {"weight_Y", "weight_U", "weight_V", NULL };
 188         double (*func2[])(void *, double, double) = { weight_Y, weight_U, weight_V, p[plane], NULL };
 189
 190         ret = av_expr_parse(&fftfilt->weight_expr[plane], fftfilt->weight_str[plane], var_names,
 191                             NULL, NULL, func2_names, func2, 0, ctx);
 192         if (ret < 0)
 193             break;
 194     }
 195     return ret;
 196 }
 197
 198 static int config_props(AVFilterLink *inlink)
 199 {
 200     FFTFILTContext *fftfilt = inlink->dst->priv;
 201     const AVPixFmtDescriptor *desc;
 202     int rdft_hbits, rdft_vbits, i, j, plane;
 203     double values[VAR_VARS_NB];
 204
 205     desc = av_pix_fmt_desc_get(inlink->format);
 206     for (i = 0; i < desc->nb_components; i++) {
 207         int w = inlink->w;
 208         int h = inlink->h;
 209
 210         /* RDFT - Array initialization for Horizontal pass*/
 211         for (rdft_hbits = 1; 1 << rdft_hbits < w*10/9; rdft_hbits++);
 212         fftfilt->rdft_hbits[i] = rdft_hbits;
 213         fftfilt->rdft_hlen[i] = 1 << rdft_hbits;
 214         if (!(fftfilt->rdft_hdata[i] = av_malloc_array(h, fftfilt->rdft_hlen[i] * sizeof(FFTSample))))
 215             return AVERROR(ENOMEM);
 216
 217         /* RDFT - Array initialization for Vertical pass*/
 218         for (rdft_vbits = 1; 1 << rdft_vbits < h*10/9; rdft_vbits++);
 219         fftfilt->rdft_vbits[i] = rdft_vbits;
 220         fftfilt->rdft_vlen[i] = 1 << rdft_vbits;
 221         if (!(fftfilt->rdft_vdata[i] = av_malloc_array(fftfilt->rdft_hlen[i], fftfilt->rdft_vlen[i] * sizeof(FFTSample))))
 222             return AVERROR(ENOMEM);
 223     }
 224
 225     /*Luminance value - Array initialization*/
 226     values[VAR_W] = inlink->w;
 227     values[VAR_H] = inlink->h;
 228     for (plane = 0; plane < 3; plane++)
 229     {
 230         if(!(fftfilt->weight[plane] = av_malloc_array(fftfilt->rdft_hlen[plane], fftfilt->rdft_vlen[plane] * sizeof(double))))
 231             return AVERROR(ENOMEM);
 232         for (i = 0; i < fftfilt->rdft_hlen[plane]; i++)
 233         {
 234             values[VAR_X] = i;
 235             for (j = 0; j < fftfilt->rdft_vlen[plane]; j++)
 236             {
 237                 values[VAR_Y] = j;
 238                 fftfilt->weight[plane][i * fftfilt->rdft_vlen[plane] + j] =
 239                 av_expr_eval(fftfilt->weight_expr[plane], values, fftfilt);
 240             }
 241         }
 242     }
 243     return 0;
 244 }
 245
 246 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 247 {
 248     AVFilterContext *ctx = inlink->dst;
 249     AVFilterLink *outlink = inlink->dst->outputs[0];
 250     const AVPixFmtDescriptor *desc;
 251     FFTFILTContext *fftfilt = ctx->priv;
 252     AVFrame *out;
 253     int i, j, plane;
 254
 255     out = ff_get_video_buffer(outlink, inlink->w, inlink->h);
 256     if (!out)
 257         return AVERROR(ENOMEM);
 258
 259     av_frame_copy_props(out, in);
 260
 261     desc = av_pix_fmt_desc_get(inlink->format);
 262     for (plane = 0; plane < desc->nb_components; plane++) {
 263         int w = inlink->w;
 264         int h = inlink->h;
 265
 266         if (plane == 1 || plane == 2) {
 267             w = FF_CEIL_RSHIFT(w, desc->log2_chroma_w);
 268             h = FF_CEIL_RSHIFT(h, desc->log2_chroma_h);
 269         }
 270
 271         rdft_horizontal(fftfilt, in, w, h, plane);
 272         rdft_vertical(fftfilt, h, plane);
 273
 274         /*Change user defined parameters*/
 275         for (i = 0; i < fftfilt->rdft_hlen[plane]; i++)
 276             for (j = 0; j < fftfilt->rdft_vlen[plane]; j++)
 277                 fftfilt->rdft_vdata[plane][i * fftfilt->rdft_vlen[plane] + j] *=
 278                   fftfilt->weight[plane][i * fftfilt->rdft_vlen[plane] + j];
 279
 280         fftfilt->rdft_vdata[plane][0] += fftfilt->rdft_hlen[plane] * fftfilt->rdft_vlen[plane] * fftfilt->dc[plane];
 281
 282         irdft_vertical(fftfilt, h, plane);
 283         irdft_horizontal(fftfilt, out, w, h, plane);
 284     }
 285
 286     av_frame_free(&in);
 287     return ff_filter_frame(outlink, out);
 288 }
 289
 290 static av_cold void uninit(AVFilterContext *ctx)
 291 {
 292     FFTFILTContext *fftfilt = ctx->priv;
 293     int i;
 294     for (i = 0; i < 3; i++) {
 295         av_free(fftfilt->rdft_hdata[i]);
 296         av_free(fftfilt->rdft_vdata[i]);
 297         av_expr_free(fftfilt->weight_expr[i]);
 298         av_free(fftfilt->weight[i]);
 299     }
 300 }
 301
 302 static int query_formats(AVFilterContext *ctx)
 303 {
 304     static const enum AVPixelFormat pixel_fmts_fftfilt[] = {
 305         AV_PIX_FMT_GRAY8,
 306         AV_PIX_FMT_YUV444P,
 307         AV_PIX_FMT_NONE
 308     };
 309
 310     AVFilterFormats *fmts_list = ff_make_format_list(pixel_fmts_fftfilt);
 311     if (!fmts_list)
 312         return AVERROR(ENOMEM);
 313     return ff_set_common_formats(ctx, fmts_list);
 314 }
 315
 316 static const AVFilterPad fftfilt_inputs[] = {
 317     {
 318         .name = "default",
 319         .type = AVMEDIA_TYPE_VIDEO,
 320         .config_props = config_props,
 321         .filter_frame = filter_frame,
 322     },
 323     { NULL }
 324 };
 325
 326 static const AVFilterPad fftfilt_outputs[] = {
 327     {
 328         .name = "default",
 329         .type = AVMEDIA_TYPE_VIDEO,
 330     },
 331     { NULL }
 332 };
 333
 334 AVFilter ff_vf_fftfilt = {
 335     .name            = "fftfilt",
 336     .description     = NULL_IF_CONFIG_SMALL("Apply arbitrary expressions to samples in frequency domain"),
 337     .priv_size       = sizeof(FFTFILTContext),
 338     .priv_class      = &fftfilt_class,
 339     .inputs          = fftfilt_inputs,
 340     .outputs         = fftfilt_outputs,
 341     .query_formats   = query_formats,
 342     .init            = initialize,
 343     .uninit          = uninit,
 344 };