avfilter: add find_rect filter
[ffmpeg.git] / libavfilter / vf_find_rect.c
1 /*
2  * Copyright (c) 2014-2015 Michael Niedermayer <michaelni@gmx.at>
3  *
4  * This file is part of FFmpeg.
5  *
6  * FFmpeg is free software; you can redistribute it and/or modify
7  * it under the terms of the GNU General Public License as published by
8  * the Free Software Foundation; either version 2 of the License, or
9  * (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
14  * GNU General Public License for more details.
15  *
16  * You should have received a copy of the GNU General Public License along
17  * with FFmpeg; if not, write to the Free Software Foundation, Inc.,
18  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
19  */
20
21 /**
22  * @todo switch to dualinput
23  */
24
25 #include "libavutil/avassert.h"
26 #include "libavutil/imgutils.h"
27 #include "libavutil/opt.h"
28 #include "internal.h"
29
30 #include "lavfutils.h"
31
32 #define MAX_MIPMAPS 5
33
34 typedef struct FOCContext {
35     AVClass *class;
36     float threshold;
37     int mipmaps;
38     int xmin, ymin, xmax, ymax;
39     char *obj_filename;
40     int last_x, last_y;
41     AVFrame *obj_frame;
42     AVFrame *needle_frame[MAX_MIPMAPS];
43     AVFrame *haystack_frame[MAX_MIPMAPS];
44 } FOCContext;
45
46 #define OFFSET(x) offsetof(FOCContext, x)
47 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
48 static const AVOption foc_options[] = {
49     { "object", "object bitmap filename", OFFSET(obj_filename), AV_OPT_TYPE_STRING, {.str = NULL}, .flags = FLAGS },
50     { "threshold", "set threshold", OFFSET(threshold), AV_OPT_TYPE_FLOAT, {.dbl = 0.5}, 0, 1.0, FLAGS },
51     { "mipmaps", "set mipmaps", OFFSET(mipmaps), AV_OPT_TYPE_INT, {.i64 = 3}, 1, MAX_MIPMAPS, FLAGS },
52     { "xmin", "", OFFSET(xmin), AV_OPT_TYPE_INT, {.i64 = 0}, 0, INT_MAX, FLAGS },
53     { "ymin", "", OFFSET(ymin), AV_OPT_TYPE_INT, {.i64 = 0}, 0, INT_MAX, FLAGS },
54     { "xmax", "", OFFSET(xmax), AV_OPT_TYPE_INT, {.i64 = 0}, 0, INT_MAX, FLAGS },
55     { "ymax", "", OFFSET(ymax), AV_OPT_TYPE_INT, {.i64 = 0}, 0, INT_MAX, FLAGS },
56     { NULL }
57 };
58
59 static const AVClass foc_class = {
60     .class_name       = "find_rect",
61     .item_name        = av_default_item_name,
62     .option           = foc_options,
63     .version          = LIBAVUTIL_VERSION_INT,
64     .category         = AV_CLASS_CATEGORY_FILTER,
65 };
66
67 static int query_formats(AVFilterContext *ctx)
68 {
69     static const enum AVPixelFormat pix_fmts[] = {
70         AV_PIX_FMT_YUV420P,
71         AV_PIX_FMT_YUVJ420P,
72         AV_PIX_FMT_NONE
73     };
74
75     return ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
76 }
77
78 static AVFrame *downscale(AVFrame *in)
79 {
80     int x, y;
81     AVFrame *frame = av_frame_alloc();
82     uint8_t *src, *dst;
83     if (!frame)
84         return NULL;
85
86     frame->format = in->format;
87     frame->width  = (in->width + 1) / 2;
88     frame->height = (in->height+ 1) / 2;
89
90     if (av_frame_get_buffer(frame, 32) < 0) {
91         av_frame_free(&frame);
92         return NULL;
93     }
94     src = in   ->data[0];
95     dst = frame->data[0];
96
97     for(y = 0; y < frame->height; y++) {
98         for(x = 0; x < frame->width; x++) {
99             dst[x] = (  src[2*x+0]
100                       + src[2*x+1]
101                       + src[2*x+0 + in->linesize[0]]
102                       + src[2*x+1 + in->linesize[0]]
103                       + 2) >> 2;
104         }
105         src += 2*in->linesize[0];
106         dst += frame->linesize[0];
107     }
108     return frame;
109 }
110
111 static float compare(const AVFrame *haystack, const AVFrame *obj, int offx, int offy)
112 {
113     int x,y;
114     int o_sum_v = 0;
115     int h_sum_v = 0;
116     int64_t oo_sum_v = 0;
117     int64_t hh_sum_v = 0;
118     int64_t oh_sum_v = 0;
119     float c;
120     int n = obj->height * obj->width;
121     const uint8_t *odat = obj     ->data[0];
122     const uint8_t *hdat = haystack->data[0] + offx + offy * haystack->linesize[0];
123     int64_t o_sigma, h_sigma;
124
125     for(y = 0; y < obj->height; y++) {
126         for(x = 0; x < obj->width; x++) {
127             int o_v = odat[x];
128             int h_v = hdat[x];
129             o_sum_v += o_v;
130             h_sum_v += h_v;
131             oo_sum_v += o_v * o_v;
132             hh_sum_v += h_v * h_v;
133             oh_sum_v += o_v * h_v;
134         }
135         odat += obj->linesize[0];
136         hdat += haystack->linesize[0];
137     }
138     o_sigma = n*oo_sum_v - o_sum_v*(int64_t)o_sum_v;
139     h_sigma = n*hh_sum_v - h_sum_v*(int64_t)h_sum_v;
140
141     if (o_sigma == 0 || h_sigma == 0)
142         return 1.0;
143
144     c = (n*oh_sum_v - o_sum_v*(int64_t)h_sum_v) / (sqrt(o_sigma)*sqrt(h_sigma));
145
146     return 1 - fabs(c);
147 }
148
149 static int config_input(AVFilterLink *inlink)
150 {
151     AVFilterContext *ctx = inlink->dst;
152     FOCContext *foc = ctx->priv;
153
154     if (foc->xmax <= 0)
155         foc->xmax = inlink->w - foc->obj_frame->width;
156     if (foc->ymax <= 0)
157         foc->ymax = inlink->h - foc->obj_frame->height;
158
159     return 0;
160 }
161
162 static float search(FOCContext *foc, int pass, int maxpass, int xmin, int xmax, int ymin, int ymax, int *best_x, int *best_y, float best_score)
163 {
164     int x, y;
165
166     if (pass + 1 <= maxpass) {
167         int sub_x, sub_y;
168         search(foc, pass+1, maxpass, xmin>>1, (xmax+1)>>1, ymin>>1, (ymax+1)>>1, &sub_x, &sub_y, 1.0);
169         xmin = FFMAX(xmin, 2*sub_x - 4);
170         xmax = FFMIN(xmax, 2*sub_x + 4);
171         ymin = FFMAX(ymin, 2*sub_y - 4);
172         ymax = FFMIN(ymax, 2*sub_y + 4);
173     }
174
175     for (y = ymin; y <= ymax; y++) {
176         for (x = xmin; x <= xmax; x++) {
177             float score = compare(foc->haystack_frame[pass], foc->needle_frame[pass], x, y);
178             av_assert0(score != 0);
179             if (score < best_score) {
180                 best_score = score;
181                 *best_x = x;
182                 *best_y = y;
183             }
184         }
185     }
186     return best_score;
187 }
188
189 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
190 {
191     AVFilterContext *ctx = inlink->dst;
192     FOCContext *foc = ctx->priv;
193     float best_score;
194     int best_x, best_y;
195     int i;
196
197     foc->haystack_frame[0] = av_frame_clone(in);
198     for (i=1; i<foc->mipmaps; i++) {
199         foc->haystack_frame[i] = downscale(foc->haystack_frame[i-1]);
200     }
201
202     best_score = search(foc, 0, 0,
203                         FFMAX(foc->xmin, foc->last_x - 8),
204                         FFMIN(foc->xmax, foc->last_x + 8),
205                         FFMAX(foc->ymin, foc->last_y - 8),
206                         FFMIN(foc->ymax, foc->last_y + 8),
207                         &best_x, &best_y, 1.0);
208
209     best_score = search(foc, 0, foc->mipmaps - 1, foc->xmin, foc->xmax, foc->ymin, foc->ymax,
210                         &best_x, &best_y, best_score);
211
212     for (i=0; i<MAX_MIPMAPS; i++) {
213         av_frame_free(&foc->haystack_frame[i]);
214     }
215
216     if (best_score > foc->threshold) {
217         return ff_filter_frame(ctx->outputs[0], in);
218     }
219
220     av_log(ctx, AV_LOG_DEBUG, "Found at %d %d score %f\n", best_x, best_y, best_score);
221     foc->last_x = best_x;
222     foc->last_y = best_y;
223
224     av_frame_make_writable(in);
225
226     av_dict_set_int(&in->metadata, "lavfi.rect.w", foc->obj_frame->width, 0);
227     av_dict_set_int(&in->metadata, "lavfi.rect.h", foc->obj_frame->height, 0);
228     av_dict_set_int(&in->metadata, "lavfi.rect.x", best_x, 0);
229     av_dict_set_int(&in->metadata, "lavfi.rect.y", best_y, 0);
230
231     return ff_filter_frame(ctx->outputs[0], in);
232 }
233
234 static av_cold void uninit(AVFilterContext *ctx)
235 {
236     FOCContext *foc = ctx->priv;
237     int i;
238
239     for (i = 0; i < MAX_MIPMAPS; i++) {
240         av_frame_free(&foc->needle_frame[i]);
241         av_frame_free(&foc->haystack_frame[i]);
242     }
243
244     if (foc->obj_frame)
245         av_freep(&foc->obj_frame->data[0]);
246     av_frame_free(&foc->obj_frame);
247 }
248
249 static av_cold int init(AVFilterContext *ctx)
250 {
251     FOCContext *foc = ctx->priv;
252     int ret, i;
253
254     if (!foc->obj_filename) {
255         av_log(ctx, AV_LOG_ERROR, "object filename not set\n");
256         return AVERROR(EINVAL);
257     }
258
259     foc->obj_frame = av_frame_alloc();
260     if (!foc->obj_frame)
261         return AVERROR(ENOMEM);
262
263     if ((ret = ff_load_image(foc->obj_frame->data, foc->obj_frame->linesize,
264                              &foc->obj_frame->width, &foc->obj_frame->height,
265                              &foc->obj_frame->format, foc->obj_filename, ctx)) < 0)
266         return ret;
267
268     if (foc->obj_frame->format != AV_PIX_FMT_GRAY8) {
269         av_log(ctx, AV_LOG_ERROR, "object image is not a grayscale image\n");
270         return AVERROR(EINVAL);
271     }
272
273     foc->needle_frame[0] = av_frame_clone(foc->obj_frame);
274     for (i = 1; i < foc->mipmaps; i++) {
275         foc->needle_frame[i] = downscale(foc->needle_frame[i-1]);
276         if (!foc->needle_frame[i])
277             return AVERROR(ENOMEM);
278     }
279
280     return 0;
281 }
282
283 static const AVFilterPad foc_inputs[] = {
284     {
285         .name         = "default",
286         .type         = AVMEDIA_TYPE_VIDEO,
287         .config_props = config_input,
288         .filter_frame = filter_frame,
289     },
290     { NULL }
291 };
292
293 static const AVFilterPad foc_outputs[] = {
294     {
295         .name = "default",
296         .type = AVMEDIA_TYPE_VIDEO,
297     },
298     { NULL }
299 };
300
301 AVFilter ff_vf_find_rect = {
302     .name            = "find_rect",
303     .description     = NULL_IF_CONFIG_SMALL("Find a user specified object"),
304     .priv_size       = sizeof(FOCContext),
305     .init            = init,
306     .uninit          = uninit,
307     .query_formats   = query_formats,
308     .inputs          = foc_inputs,
309     .outputs         = foc_outputs,
310     .priv_class      = &foc_class,
311 };