blob: 307b41a7338a4d1b5fceebffcc981452a7628be3
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 { |
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 | }; |
347 |