avfilter/vf_despill: fix assigment
[ffmpeg.git] / libavfilter / vf_convolve.c
1 /*
2  * Copyright (c) 2017 Paul B Mahol
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU Lesser General Public
8  * License as published by the Free Software Foundation; either
9  * version 2.1 of the License, 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 #include "libavutil/imgutils.h"
22 #include "libavutil/opt.h"
23 #include "libavutil/pixdesc.h"
24 #include "libavcodec/avfft.h"
25
26 #include "avfilter.h"
27 #include "formats.h"
28 #include "framesync.h"
29 #include "internal.h"
30 #include "video.h"
31
32 typedef struct ConvolveContext {
33     const AVClass *class;
34     FFFrameSync fs;
35
36     FFTContext *fft[4];
37     FFTContext *ifft[4];
38
39     int fft_bits[4];
40     int fft_len[4];
41     int planewidth[4];
42     int planeheight[4];
43
44     FFTComplex *fft_hdata[4];
45     FFTComplex *fft_vdata[4];
46     FFTComplex *fft_hdata_impulse[4];
47     FFTComplex *fft_vdata_impulse[4];
48
49     int depth;
50     int planes;
51     int impulse;
52     int nb_planes;
53     int got_impulse[4];
54 } ConvolveContext;
55
56 #define OFFSET(x) offsetof(ConvolveContext, x)
57 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
58
59 static const AVOption convolve_options[] = {
60     { "planes",  "set planes to convolve",                  OFFSET(planes),   AV_OPT_TYPE_INT,   {.i64=7}, 0, 15, FLAGS },
61     { "impulse", "when to process impulses",                OFFSET(impulse),  AV_OPT_TYPE_INT,   {.i64=1}, 0,  1, FLAGS, "impulse" },
62     {   "first", "process only first impulse, ignore rest", 0,                AV_OPT_TYPE_CONST, {.i64=0}, 0,  0, FLAGS, "impulse" },
63     {   "all",   "process all impulses",                    0,                AV_OPT_TYPE_CONST, {.i64=1}, 0,  0, FLAGS, "impulse" },
64     { NULL },
65 };
66
67 FRAMESYNC_DEFINE_CLASS(convolve, ConvolveContext, fs);
68
69 static int query_formats(AVFilterContext *ctx)
70 {
71     static const enum AVPixelFormat pixel_fmts_fftfilt[] = {
72         AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV440P,
73         AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P,
74         AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUV420P,
75         AV_PIX_FMT_YUVJ422P, AV_PIX_FMT_YUVJ420P,
76         AV_PIX_FMT_YUVJ411P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_YUV410P,
77         AV_PIX_FMT_YUV420P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV444P9,
78         AV_PIX_FMT_YUV420P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV444P10,
79         AV_PIX_FMT_YUV420P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV440P12,
80         AV_PIX_FMT_YUV420P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV444P14,
81         AV_PIX_FMT_YUV420P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV444P16,
82         AV_PIX_FMT_YUVA420P9, AV_PIX_FMT_YUVA422P9, AV_PIX_FMT_YUVA444P9,
83         AV_PIX_FMT_YUVA420P10, AV_PIX_FMT_YUVA422P10, AV_PIX_FMT_YUVA444P10,
84         AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16,
85         AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10,
86         AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16,
87         AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRAP10, AV_PIX_FMT_GBRAP12, AV_PIX_FMT_GBRAP16,
88         AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY9, AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY16,
89         AV_PIX_FMT_NONE
90     };
91
92     AVFilterFormats *fmts_list = ff_make_format_list(pixel_fmts_fftfilt);
93     if (!fmts_list)
94         return AVERROR(ENOMEM);
95     return ff_set_common_formats(ctx, fmts_list);
96 }
97
98 static int config_input_main(AVFilterLink *inlink)
99 {
100     ConvolveContext *s = inlink->dst->priv;
101     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
102     int fft_bits, i;
103
104     s->planewidth[1] = s->planewidth[2] = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
105     s->planewidth[0] = s->planewidth[3] = inlink->w;
106     s->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
107     s->planeheight[0] = s->planeheight[3] = inlink->h;
108
109     s->nb_planes = desc->nb_components;
110     s->depth = desc->comp[0].depth;
111
112     for (i = 0; i < s->nb_planes; i++) {
113         int w = s->planewidth[i];
114         int h = s->planeheight[i];
115         int n = FFMAX(w, h) * 10/9;
116
117         for (fft_bits = 1; 1 << fft_bits < n; fft_bits++);
118
119         s->fft_bits[i] = fft_bits;
120         s->fft_len[i] = 1 << s->fft_bits[i];
121
122         if (!(s->fft_hdata[i] = av_calloc(s->fft_len[i], s->fft_len[i] * sizeof(FFTComplex))))
123             return AVERROR(ENOMEM);
124
125         if (!(s->fft_vdata[i] = av_calloc(s->fft_len[i], s->fft_len[i] * sizeof(FFTComplex))))
126             return AVERROR(ENOMEM);
127
128         if (!(s->fft_hdata_impulse[i] = av_calloc(s->fft_len[i], s->fft_len[i] * sizeof(FFTComplex))))
129             return AVERROR(ENOMEM);
130
131         if (!(s->fft_vdata_impulse[i] = av_calloc(s->fft_len[i], s->fft_len[i] * sizeof(FFTComplex))))
132             return AVERROR(ENOMEM);
133     }
134
135     return 0;
136 }
137
138 static int config_input_impulse(AVFilterLink *inlink)
139 {
140     AVFilterContext *ctx  = inlink->dst;
141
142     if (ctx->inputs[0]->w != ctx->inputs[1]->w ||
143         ctx->inputs[0]->h != ctx->inputs[1]->h) {
144         av_log(ctx, AV_LOG_ERROR, "Width and height of input videos must be same.\n");
145         return AVERROR(EINVAL);
146     }
147     if (ctx->inputs[0]->format != ctx->inputs[1]->format) {
148         av_log(ctx, AV_LOG_ERROR, "Inputs must be of same pixel format.\n");
149         return AVERROR(EINVAL);
150     }
151
152     return 0;
153 }
154
155 static void fft_horizontal(ConvolveContext *s, FFTComplex *fft_hdata,
156                            AVFrame *in, int w, int h, int n, int plane, float scale)
157 {
158     int y, x;
159
160     for (y = 0; y < h; y++) {
161         if (s->depth == 8) {
162             const uint8_t *src = in->data[plane] + in->linesize[plane] * y;
163
164             for (x = 0; x < w; x++) {
165                 fft_hdata[y * n + x].re = src[x] * scale;
166                 fft_hdata[y * n + x].im = 0;
167             }
168         } else {
169             const uint16_t *src = (const uint16_t *)(in->data[plane] + in->linesize[plane] * y);
170
171             for (x = 0; x < w; x++) {
172                 fft_hdata[y * n + x].re = src[x] * scale;
173                 fft_hdata[y * n + x].im = 0;
174             }
175         }
176         for (; x < n; x++) {
177             fft_hdata[y * n + x].re = 0;
178             fft_hdata[y * n + x].im = 0;
179         }
180     }
181
182     for (; y < n; y++) {
183         for (x = 0; x < n; x++) {
184             fft_hdata[y * n + x].re = 0;
185             fft_hdata[y * n + x].im = 0;
186         }
187     }
188
189     for (y = 0; y < n; y++) {
190         av_fft_permute(s->fft[plane], fft_hdata + y * n);
191         av_fft_calc(s->fft[plane], fft_hdata + y * n);
192     }
193 }
194
195 static void fft_vertical(ConvolveContext *s, FFTComplex *fft_hdata, FFTComplex *fft_vdata,
196                          int n, int plane)
197 {
198     int y, x;
199
200     for (y = 0; y < n; y++) {
201         for (x = 0; x < n; x++) {
202             fft_vdata[y * n + x].re = fft_hdata[x * n + y].re;
203             fft_vdata[y * n + x].im = fft_hdata[x * n + y].im;
204         }
205         for (; x < n; x++) {
206             fft_vdata[y * n + x].re = 0;
207             fft_vdata[y * n + x].im = 0;
208         }
209         av_fft_permute(s->fft[plane], fft_vdata + y * n);
210         av_fft_calc(s->fft[plane], fft_vdata + y * n);
211     }
212 }
213
214 static void ifft_vertical(ConvolveContext *s, int n, int plane)
215 {
216     int y, x;
217
218     for (y = 0; y < n; y++) {
219         av_fft_permute(s->ifft[plane], s->fft_vdata[plane] + y * n);
220         av_fft_calc(s->ifft[plane], s->fft_vdata[plane] + y * n);
221         for (x = 0; x < n; x++) {
222             s->fft_hdata[plane][x * n + y].re = s->fft_vdata[plane][y * n + x].re;
223             s->fft_hdata[plane][x * n + y].im = s->fft_vdata[plane][y * n + x].im;
224         }
225     }
226 }
227
228 static void ifft_horizontal(ConvolveContext *s, AVFrame *out,
229                             int w, int h, int n, int plane)
230 {
231     const float scale = 1.f / (n * n);
232     const int max = (1 << s->depth) - 1;
233     const int oh = h / 2;
234     const int ow = w / 2;
235     int y, x;
236
237     for (y = 0; y < n; y++) {
238         av_fft_permute(s->ifft[plane], s->fft_hdata[plane] + y * n);
239         av_fft_calc(s->ifft[plane], s->fft_hdata[plane] + y * n);
240     }
241
242     if (s->depth == 8) {
243         for (y = 0; y < h; y++) {
244             uint8_t *dst = out->data[plane] + y * out->linesize[plane];
245             for (x = 0; x < w; x++)
246                 dst[x] = av_clip_uint8(s->fft_hdata[plane][(y+oh) * n + x+ow].re * scale);
247         }
248     } else {
249         for (y = 0; y < h; y++) {
250             uint16_t *dst = (uint16_t *)(out->data[plane] + y * out->linesize[plane]);
251             for (x = 0; x < w; x++)
252                 dst[x] = av_clip(s->fft_hdata[plane][(y+oh) * n + x+ow].re * scale, 0, max);
253         }
254     }
255 }
256
257 static int do_convolve(FFFrameSync *fs)
258 {
259     AVFilterContext *ctx = fs->parent;
260     AVFilterLink *outlink = ctx->outputs[0];
261     ConvolveContext *s = ctx->priv;
262     AVFrame *mainpic = NULL, *impulsepic = NULL;
263     int ret, y, x, plane;
264
265     ret = ff_framesync_dualinput_get(fs, &mainpic, &impulsepic);
266     if (ret < 0)
267         return ret;
268     if (!impulsepic)
269         return ff_filter_frame(outlink, mainpic);
270
271     for (plane = 0; plane < s->nb_planes; plane++) {
272         const int n = s->fft_len[plane];
273         const int w = s->planewidth[plane];
274         const int h = s->planeheight[plane];
275         float total = 0;
276
277         if (!(s->planes & (1 << plane))) {
278             continue;
279         }
280
281         fft_horizontal(s, s->fft_hdata[plane], mainpic, w, h, n, plane, 1.f);
282         fft_vertical(s, s->fft_hdata[plane], s->fft_vdata[plane],
283                      n, plane);
284
285         if ((!s->impulse && !s->got_impulse[plane]) || s->impulse) {
286             if (s->depth == 8) {
287                 for (y = 0; y < h; y++) {
288                     const uint8_t *src = (const uint8_t *)(impulsepic->data[plane] + y * impulsepic->linesize[plane]) ;
289                     for (x = 0; x < w; x++) {
290                         total += src[x];
291                     }
292                 }
293             } else {
294                 for (y = 0; y < h; y++) {
295                     const uint16_t *src = (const uint16_t *)(impulsepic->data[plane] + y * impulsepic->linesize[plane]) ;
296                     for (x = 0; x < w; x++) {
297                         total += src[x];
298                     }
299                 }
300             }
301             total = FFMAX(1, total);
302
303             fft_horizontal(s, s->fft_hdata_impulse[plane], impulsepic, w, h, n, plane, 1 / total);
304             fft_vertical(s, s->fft_hdata_impulse[plane], s->fft_vdata_impulse[plane],
305                          n, plane);
306
307             s->got_impulse[plane] = 1;
308         }
309
310         for (y = 0; y < n; y++) {
311             for (x = 0; x < n; x++) {
312                 FFTSample re, im, ire, iim;
313
314                 re = s->fft_vdata[plane][y*n + x].re;
315                 im = s->fft_vdata[plane][y*n + x].im;
316                 ire = s->fft_vdata_impulse[plane][y*n + x].re;
317                 iim = s->fft_vdata_impulse[plane][y*n + x].im;
318
319                 s->fft_vdata[plane][y*n + x].re = ire * re - iim * im;
320                 s->fft_vdata[plane][y*n + x].im = iim * re + ire * im;
321             }
322         }
323
324         ifft_vertical(s, n, plane);
325         ifft_horizontal(s, mainpic, w, h, n, plane);
326     }
327
328     return ff_filter_frame(outlink, mainpic);
329 }
330
331 static int config_output(AVFilterLink *outlink)
332 {
333     AVFilterContext *ctx = outlink->src;
334     ConvolveContext *s = ctx->priv;
335     AVFilterLink *mainlink = ctx->inputs[0];
336     int ret, i;
337
338     s->fs.on_event = do_convolve;
339     ret = ff_framesync_init_dualinput(&s->fs, ctx);
340     if (ret < 0)
341         return ret;
342     outlink->w = mainlink->w;
343     outlink->h = mainlink->h;
344     outlink->time_base = mainlink->time_base;
345     outlink->sample_aspect_ratio = mainlink->sample_aspect_ratio;
346     outlink->frame_rate = mainlink->frame_rate;
347
348     if ((ret = ff_framesync_configure(&s->fs)) < 0)
349         return ret;
350
351     for (i = 0; i < s->nb_planes; i++) {
352         s->fft[i]  = av_fft_init(s->fft_bits[i], 0);
353         s->ifft[i] = av_fft_init(s->fft_bits[i], 1);
354         if (!s->fft[i] || !s->ifft[i])
355             return AVERROR(ENOMEM);
356     }
357
358     return 0;
359 }
360
361 static int activate(AVFilterContext *ctx)
362 {
363     ConvolveContext *s = ctx->priv;
364     return ff_framesync_activate(&s->fs);
365 }
366
367 static av_cold void uninit(AVFilterContext *ctx)
368 {
369     ConvolveContext *s = ctx->priv;
370     int i;
371
372     for (i = 0; i < 4; i++) {
373         av_freep(&s->fft_hdata[i]);
374         av_freep(&s->fft_vdata[i]);
375         av_freep(&s->fft_hdata_impulse[i]);
376         av_freep(&s->fft_vdata_impulse[i]);
377         av_fft_end(s->fft[i]);
378         av_fft_end(s->ifft[i]);
379     }
380
381     ff_framesync_uninit(&s->fs);
382 }
383
384 static const AVFilterPad convolve_inputs[] = {
385     {
386         .name          = "main",
387         .type          = AVMEDIA_TYPE_VIDEO,
388         .config_props  = config_input_main,
389     },{
390         .name          = "impulse",
391         .type          = AVMEDIA_TYPE_VIDEO,
392         .config_props  = config_input_impulse,
393     },
394     { NULL }
395 };
396
397 static const AVFilterPad convolve_outputs[] = {
398     {
399         .name          = "default",
400         .type          = AVMEDIA_TYPE_VIDEO,
401         .config_props  = config_output,
402     },
403     { NULL }
404 };
405
406 AVFilter ff_vf_convolve = {
407     .name          = "convolve",
408     .description   = NULL_IF_CONFIG_SMALL("Convolve first video stream with second video stream."),
409     .preinit       = convolve_framesync_preinit,
410     .uninit        = uninit,
411     .query_formats = query_formats,
412     .activate      = activate,
413     .priv_size     = sizeof(ConvolveContext),
414     .priv_class    = &convolve_class,
415     .inputs        = convolve_inputs,
416     .outputs       = convolve_outputs,
417     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL,
418 };