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/imgutils.h"
26 #include "libavutil/opt.h"
27 #include "internal.h"
28
29 #include "lavfutils.h"
30
31 #define MAX_MIPMAPS 5
32
33 typedef struct FOCContext {
34 AVClass *class;
35 float threshold;
36 int mipmaps;
37 int xmin, ymin, xmax, ymax;
38 char *obj_filename;
39 int last_x, last_y;
40 AVFrame *obj_frame;
41 AVFrame *needle_frame[MAX_MIPMAPS];
42 AVFrame *haystack_frame[MAX_MIPMAPS];
43 int discard;
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 find_rect_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 { "discard", "", OFFSET(discard), AV_OPT_TYPE_BOOL, {.i64 = 0}, 0, 1, FLAGS },
57 { NULL }
58 };
59
60 AVFILTER_DEFINE_CLASS(find_rect);
61
downscale(AVFrame * in)62 static AVFrame *downscale(AVFrame *in)
63 {
64 int x, y;
65 AVFrame *frame = av_frame_alloc();
66 uint8_t *src, *dst;
67 if (!frame)
68 return NULL;
69
70 frame->format = in->format;
71 frame->width = (in->width + 1) / 2;
72 frame->height = (in->height+ 1) / 2;
73
74 if (av_frame_get_buffer(frame, 0) < 0) {
75 av_frame_free(&frame);
76 return NULL;
77 }
78 src = in ->data[0];
79 dst = frame->data[0];
80
81 for(y = 0; y < frame->height; y++) {
82 for(x = 0; x < frame->width; x++) {
83 dst[x] = ( src[2*x+0]
84 + src[2*x+1]
85 + src[2*x+0 + in->linesize[0]]
86 + src[2*x+1 + in->linesize[0]]
87 + 2) >> 2;
88 }
89 src += 2*in->linesize[0];
90 dst += frame->linesize[0];
91 }
92 return frame;
93 }
94
compare(const AVFrame * haystack,const AVFrame * obj,int offx,int offy)95 static float compare(const AVFrame *haystack, const AVFrame *obj, int offx, int offy)
96 {
97 int x,y;
98 int o_sum_v = 0;
99 int h_sum_v = 0;
100 int64_t oo_sum_v = 0;
101 int64_t hh_sum_v = 0;
102 int64_t oh_sum_v = 0;
103 float c;
104 int n = obj->height * obj->width;
105 const uint8_t *odat = obj ->data[0];
106 const uint8_t *hdat = haystack->data[0] + offx + offy * haystack->linesize[0];
107 int64_t o_sigma, h_sigma;
108
109 for(y = 0; y < obj->height; y++) {
110 for(x = 0; x < obj->width; x++) {
111 int o_v = odat[x];
112 int h_v = hdat[x];
113 o_sum_v += o_v;
114 h_sum_v += h_v;
115 oo_sum_v += o_v * o_v;
116 hh_sum_v += h_v * h_v;
117 oh_sum_v += o_v * h_v;
118 }
119 odat += obj->linesize[0];
120 hdat += haystack->linesize[0];
121 }
122 o_sigma = n*oo_sum_v - o_sum_v*(int64_t)o_sum_v;
123 h_sigma = n*hh_sum_v - h_sum_v*(int64_t)h_sum_v;
124
125 if (o_sigma == 0 || h_sigma == 0)
126 return 1.0;
127
128 c = (n*oh_sum_v - o_sum_v*(int64_t)h_sum_v) / (sqrt(o_sigma)*sqrt(h_sigma));
129
130 return 1 - fabs(c);
131 }
132
config_input(AVFilterLink * inlink)133 static int config_input(AVFilterLink *inlink)
134 {
135 AVFilterContext *ctx = inlink->dst;
136 FOCContext *foc = ctx->priv;
137
138 if (foc->xmax <= 0)
139 foc->xmax = inlink->w - foc->obj_frame->width;
140 if (foc->ymax <= 0)
141 foc->ymax = inlink->h - foc->obj_frame->height;
142
143 return 0;
144 }
145
search(FOCContext * foc,int pass,int maxpass,int xmin,int xmax,int ymin,int ymax,int * best_x,int * best_y,float best_score)146 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)
147 {
148 int x, y;
149
150 if (pass + 1 <= maxpass) {
151 int sub_x, sub_y;
152 search(foc, pass+1, maxpass, xmin>>1, (xmax+1)>>1, ymin>>1, (ymax+1)>>1, &sub_x, &sub_y, 2.0);
153 xmin = FFMAX(xmin, 2*sub_x - 4);
154 xmax = FFMIN(xmax, 2*sub_x + 4);
155 ymin = FFMAX(ymin, 2*sub_y - 4);
156 ymax = FFMIN(ymax, 2*sub_y + 4);
157 }
158
159 for (y = ymin; y <= ymax; y++) {
160 for (x = xmin; x <= xmax; x++) {
161 float score = compare(foc->haystack_frame[pass], foc->needle_frame[pass], x, y);
162 if (score < best_score) {
163 best_score = score;
164 *best_x = x;
165 *best_y = y;
166 }
167 }
168 }
169 return best_score;
170 }
171
filter_frame(AVFilterLink * inlink,AVFrame * in)172 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
173 {
174 AVFilterContext *ctx = inlink->dst;
175 FOCContext *foc = ctx->priv;
176 float best_score;
177 int best_x, best_y;
178 int i;
179 char buf[32];
180
181 foc->haystack_frame[0] = av_frame_clone(in);
182 for (i=1; i<foc->mipmaps; i++) {
183 foc->haystack_frame[i] = downscale(foc->haystack_frame[i-1]);
184 }
185
186 best_score = search(foc, 0, 0,
187 FFMAX(foc->xmin, foc->last_x - 8),
188 FFMIN(foc->xmax, foc->last_x + 8),
189 FFMAX(foc->ymin, foc->last_y - 8),
190 FFMIN(foc->ymax, foc->last_y + 8),
191 &best_x, &best_y, 2.0);
192
193 best_score = search(foc, 0, foc->mipmaps - 1, foc->xmin, foc->xmax, foc->ymin, foc->ymax,
194 &best_x, &best_y, best_score);
195
196 for (i=0; i<MAX_MIPMAPS; i++) {
197 av_frame_free(&foc->haystack_frame[i]);
198 }
199
200 if (best_score > foc->threshold) {
201 if (foc->discard) {
202 av_frame_free(&in);
203 return 0;
204 } else {
205 return ff_filter_frame(ctx->outputs[0], in);
206 }
207 }
208
209 av_log(ctx, AV_LOG_INFO, "Found at n=%"PRId64" pts_time=%f x=%d y=%d with score=%f\n",
210 inlink->frame_count_out, TS2D(in->pts) * av_q2d(inlink->time_base),
211 best_x, best_y, best_score);
212 foc->last_x = best_x;
213 foc->last_y = best_y;
214
215 snprintf(buf, sizeof(buf), "%f", best_score);
216
217 av_dict_set_int(&in->metadata, "lavfi.rect.w", foc->obj_frame->width, 0);
218 av_dict_set_int(&in->metadata, "lavfi.rect.h", foc->obj_frame->height, 0);
219 av_dict_set_int(&in->metadata, "lavfi.rect.x", best_x, 0);
220 av_dict_set_int(&in->metadata, "lavfi.rect.y", best_y, 0);
221 av_dict_set(&in->metadata, "lavfi.rect.score", buf, 0);
222
223 return ff_filter_frame(ctx->outputs[0], in);
224 }
225
uninit(AVFilterContext * ctx)226 static av_cold void uninit(AVFilterContext *ctx)
227 {
228 FOCContext *foc = ctx->priv;
229 int i;
230
231 for (i = 0; i < MAX_MIPMAPS; i++) {
232 av_frame_free(&foc->needle_frame[i]);
233 av_frame_free(&foc->haystack_frame[i]);
234 }
235
236 if (foc->obj_frame)
237 av_freep(&foc->obj_frame->data[0]);
238 av_frame_free(&foc->obj_frame);
239 }
240
init(AVFilterContext * ctx)241 static av_cold int init(AVFilterContext *ctx)
242 {
243 FOCContext *foc = ctx->priv;
244 int ret, i;
245
246 if (!foc->obj_filename) {
247 av_log(ctx, AV_LOG_ERROR, "object filename not set\n");
248 return AVERROR(EINVAL);
249 }
250
251 foc->obj_frame = av_frame_alloc();
252 if (!foc->obj_frame)
253 return AVERROR(ENOMEM);
254
255 if ((ret = ff_load_image(foc->obj_frame->data, foc->obj_frame->linesize,
256 &foc->obj_frame->width, &foc->obj_frame->height,
257 &foc->obj_frame->format, foc->obj_filename, ctx)) < 0)
258 return ret;
259
260 if (foc->obj_frame->format != AV_PIX_FMT_GRAY8) {
261 av_log(ctx, AV_LOG_ERROR, "object image is not a grayscale image\n");
262 return AVERROR(EINVAL);
263 }
264
265 foc->needle_frame[0] = av_frame_clone(foc->obj_frame);
266 for (i = 1; i < foc->mipmaps; i++) {
267 foc->needle_frame[i] = downscale(foc->needle_frame[i-1]);
268 if (!foc->needle_frame[i])
269 return AVERROR(ENOMEM);
270 }
271
272 return 0;
273 }
274
275 static const AVFilterPad foc_inputs[] = {
276 {
277 .name = "default",
278 .type = AVMEDIA_TYPE_VIDEO,
279 .config_props = config_input,
280 .filter_frame = filter_frame,
281 },
282 };
283
284 static const AVFilterPad foc_outputs[] = {
285 {
286 .name = "default",
287 .type = AVMEDIA_TYPE_VIDEO,
288 },
289 };
290
291 const AVFilter ff_vf_find_rect = {
292 .name = "find_rect",
293 .description = NULL_IF_CONFIG_SMALL("Find a user specified object."),
294 .priv_size = sizeof(FOCContext),
295 .init = init,
296 .uninit = uninit,
297 .flags = AVFILTER_FLAG_METADATA_ONLY,
298 FILTER_INPUTS(foc_inputs),
299 FILTER_OUTPUTS(foc_outputs),
300 FILTER_PIXFMTS(AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUVJ420P),
301 .priv_class = &find_rect_class,
302 };
303