blob: 7ecfe7da6038c02483e774c59c01f84c4d1881a1
1 | /** |
2 | * Copyright (c) 2016 Davinder Singh (DSM_) <ds.mudhar<@gmail.com> |
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 "motion_estimation.h" |
22 | #include "libavcodec/mathops.h" |
23 | #include "libavutil/avassert.h" |
24 | #include "libavutil/common.h" |
25 | #include "libavutil/imgutils.h" |
26 | #include "libavutil/opt.h" |
27 | #include "libavutil/pixdesc.h" |
28 | #include "libavutil/motion_vector.h" |
29 | #include "avfilter.h" |
30 | #include "formats.h" |
31 | #include "internal.h" |
32 | #include "video.h" |
33 | |
34 | typedef struct MEContext { |
35 | const AVClass *class; |
36 | AVMotionEstContext me_ctx; |
37 | int method; ///< motion estimation method |
38 | |
39 | int mb_size; ///< macroblock size |
40 | int search_param; ///< search parameter |
41 | int b_width, b_height, b_count; |
42 | int log2_mb_size; |
43 | |
44 | AVFrame *prev, *cur, *next; |
45 | |
46 | int (*mv_table[3])[2][2]; ///< motion vectors of current & prev 2 frames |
47 | } MEContext; |
48 | |
49 | #define OFFSET(x) offsetof(MEContext, x) |
50 | #define FLAGS AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_FILTERING_PARAM |
51 | #define CONST(name, help, val, unit) { name, help, 0, AV_OPT_TYPE_CONST, {.i64=val}, 0, 0, FLAGS, unit } |
52 | |
53 | static const AVOption mestimate_options[] = { |
54 | { "method", "motion estimation method", OFFSET(method), AV_OPT_TYPE_INT, {.i64 = AV_ME_METHOD_ESA}, AV_ME_METHOD_ESA, AV_ME_METHOD_UMH, FLAGS, "method" }, |
55 | CONST("esa", "exhaustive search", AV_ME_METHOD_ESA, "method"), |
56 | CONST("tss", "three step search", AV_ME_METHOD_TSS, "method"), |
57 | CONST("tdls", "two dimensional logarithmic search", AV_ME_METHOD_TDLS, "method"), |
58 | CONST("ntss", "new three step search", AV_ME_METHOD_NTSS, "method"), |
59 | CONST("fss", "four step search", AV_ME_METHOD_FSS, "method"), |
60 | CONST("ds", "diamond search", AV_ME_METHOD_DS, "method"), |
61 | CONST("hexbs", "hexagon-based search", AV_ME_METHOD_HEXBS, "method"), |
62 | CONST("epzs", "enhanced predictive zonal search", AV_ME_METHOD_EPZS, "method"), |
63 | CONST("umh", "uneven multi-hexagon search", AV_ME_METHOD_UMH, "method"), |
64 | { "mb_size", "macroblock size", OFFSET(mb_size), AV_OPT_TYPE_INT, {.i64 = 16}, 8, INT_MAX, FLAGS }, |
65 | { "search_param", "search parameter", OFFSET(search_param), AV_OPT_TYPE_INT, {.i64 = 7}, 4, INT_MAX, FLAGS }, |
66 | { NULL } |
67 | }; |
68 | |
69 | AVFILTER_DEFINE_CLASS(mestimate); |
70 | |
71 | static int query_formats(AVFilterContext *ctx) |
72 | { |
73 | static const enum AVPixelFormat pix_fmts[] = { |
74 | AV_PIX_FMT_YUV410P, AV_PIX_FMT_YUV411P, |
75 | AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, |
76 | AV_PIX_FMT_YUV440P, AV_PIX_FMT_YUV444P, |
77 | AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P, |
78 | AV_PIX_FMT_YUVJ422P, AV_PIX_FMT_YUVJ420P, |
79 | AV_PIX_FMT_YUVJ411P, |
80 | AV_PIX_FMT_YUVA420P, AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUVA444P, |
81 | AV_PIX_FMT_GRAY8, |
82 | AV_PIX_FMT_NONE |
83 | }; |
84 | |
85 | AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts); |
86 | if (!fmts_list) |
87 | return AVERROR(ENOMEM); |
88 | return ff_set_common_formats(ctx, fmts_list); |
89 | } |
90 | |
91 | static int config_input(AVFilterLink *inlink) |
92 | { |
93 | MEContext *s = inlink->dst->priv; |
94 | int i; |
95 | |
96 | s->log2_mb_size = av_ceil_log2_c(s->mb_size); |
97 | s->mb_size = 1 << s->log2_mb_size; |
98 | |
99 | s->b_width = inlink->w >> s->log2_mb_size; |
100 | s->b_height = inlink->h >> s->log2_mb_size; |
101 | s->b_count = s->b_width * s->b_height; |
102 | |
103 | for (i = 0; i < 3; i++) { |
104 | s->mv_table[i] = av_mallocz_array(s->b_count, sizeof(*s->mv_table[0])); |
105 | if (!s->mv_table[i]) |
106 | return AVERROR(ENOMEM); |
107 | } |
108 | |
109 | ff_me_init_context(&s->me_ctx, s->mb_size, s->search_param, inlink->w, inlink->h, 0, (s->b_width - 1) << s->log2_mb_size, 0, (s->b_height - 1) << s->log2_mb_size); |
110 | |
111 | return 0; |
112 | } |
113 | |
114 | static void add_mv_data(AVMotionVector *mv, int mb_size, |
115 | int x, int y, int x_mv, int y_mv, int dir) |
116 | { |
117 | mv->w = mb_size; |
118 | mv->h = mb_size; |
119 | mv->dst_x = x + (mb_size >> 1); |
120 | mv->dst_y = y + (mb_size >> 1); |
121 | mv->src_x = x_mv + (mb_size >> 1); |
122 | mv->src_y = y_mv + (mb_size >> 1); |
123 | mv->source = dir ? 1 : -1; |
124 | mv->flags = 0; |
125 | } |
126 | |
127 | #define SEARCH_MV(method)\ |
128 | do {\ |
129 | for (mb_y = 0; mb_y < s->b_height; mb_y++)\ |
130 | for (mb_x = 0; mb_x < s->b_width; mb_x++) {\ |
131 | const int x_mb = mb_x << s->log2_mb_size;\ |
132 | const int y_mb = mb_y << s->log2_mb_size;\ |
133 | int mv[2] = {x_mb, y_mb};\ |
134 | ff_me_search_##method(me_ctx, x_mb, y_mb, mv);\ |
135 | add_mv_data(((AVMotionVector *) sd->data) + mv_count++, me_ctx->mb_size, x_mb, y_mb, mv[0], mv[1], dir);\ |
136 | }\ |
137 | } while (0) |
138 | |
139 | #define ADD_PRED(preds, px, py)\ |
140 | do {\ |
141 | preds.mvs[preds.nb][0] = px;\ |
142 | preds.mvs[preds.nb][1] = py;\ |
143 | preds.nb++;\ |
144 | } while(0) |
145 | |
146 | static int filter_frame(AVFilterLink *inlink, AVFrame *frame) |
147 | { |
148 | AVFilterContext *ctx = inlink->dst; |
149 | MEContext *s = ctx->priv; |
150 | AVMotionEstContext *me_ctx = &s->me_ctx; |
151 | AVFrameSideData *sd; |
152 | AVFrame *out; |
153 | int mb_x, mb_y, dir; |
154 | int32_t mv_count = 0; |
155 | int ret; |
156 | |
157 | if (frame->pts == AV_NOPTS_VALUE) { |
158 | ret = ff_filter_frame(ctx->outputs[0], frame); |
159 | return ret; |
160 | } |
161 | |
162 | av_frame_free(&s->prev); |
163 | s->prev = s->cur; |
164 | s->cur = s->next; |
165 | s->next = frame; |
166 | |
167 | s->mv_table[2] = memcpy(s->mv_table[2], s->mv_table[1], sizeof(*s->mv_table[1]) * s->b_count); |
168 | s->mv_table[1] = memcpy(s->mv_table[1], s->mv_table[0], sizeof(*s->mv_table[0]) * s->b_count); |
169 | |
170 | if (!s->cur) { |
171 | s->cur = av_frame_clone(frame); |
172 | if (!s->cur) |
173 | return AVERROR(ENOMEM); |
174 | } |
175 | |
176 | if (!s->prev) |
177 | return 0; |
178 | |
179 | out = av_frame_clone(s->cur); |
180 | if (!out) |
181 | return AVERROR(ENOMEM); |
182 | |
183 | sd = av_frame_new_side_data(out, AV_FRAME_DATA_MOTION_VECTORS, 2 * s->b_count * sizeof(AVMotionVector)); |
184 | if (!sd) { |
185 | av_frame_free(&out); |
186 | return AVERROR(ENOMEM); |
187 | } |
188 | |
189 | me_ctx->data_cur = s->cur->data[0]; |
190 | me_ctx->linesize = s->cur->linesize[0]; |
191 | |
192 | for (dir = 0; dir < 2; dir++) { |
193 | me_ctx->data_ref = (dir ? s->next : s->prev)->data[0]; |
194 | |
195 | if (s->method == AV_ME_METHOD_DS) |
196 | SEARCH_MV(ds); |
197 | else if (s->method == AV_ME_METHOD_ESA) |
198 | SEARCH_MV(esa); |
199 | else if (s->method == AV_ME_METHOD_FSS) |
200 | SEARCH_MV(fss); |
201 | else if (s->method == AV_ME_METHOD_NTSS) |
202 | SEARCH_MV(ntss); |
203 | else if (s->method == AV_ME_METHOD_TDLS) |
204 | SEARCH_MV(tdls); |
205 | else if (s->method == AV_ME_METHOD_TSS) |
206 | SEARCH_MV(tss); |
207 | else if (s->method == AV_ME_METHOD_HEXBS) |
208 | SEARCH_MV(hexbs); |
209 | else if (s->method == AV_ME_METHOD_UMH) { |
210 | for (mb_y = 0; mb_y < s->b_height; mb_y++) |
211 | for (mb_x = 0; mb_x < s->b_width; mb_x++) { |
212 | const int mb_i = mb_x + mb_y * s->b_width; |
213 | const int x_mb = mb_x << s->log2_mb_size; |
214 | const int y_mb = mb_y << s->log2_mb_size; |
215 | int mv[2] = {x_mb, y_mb}; |
216 | |
217 | AVMotionEstPredictor *preds = me_ctx->preds; |
218 | preds[0].nb = 0; |
219 | |
220 | ADD_PRED(preds[0], 0, 0); |
221 | |
222 | //left mb in current frame |
223 | if (mb_x > 0) |
224 | ADD_PRED(preds[0], s->mv_table[0][mb_i - 1][dir][0], s->mv_table[0][mb_i - 1][dir][1]); |
225 | |
226 | if (mb_y > 0) { |
227 | //top mb in current frame |
228 | ADD_PRED(preds[0], s->mv_table[0][mb_i - s->b_width][dir][0], s->mv_table[0][mb_i - s->b_width][dir][1]); |
229 | |
230 | //top-right mb in current frame |
231 | if (mb_x + 1 < s->b_width) |
232 | ADD_PRED(preds[0], s->mv_table[0][mb_i - s->b_width + 1][dir][0], s->mv_table[0][mb_i - s->b_width + 1][dir][1]); |
233 | //top-left mb in current frame |
234 | else if (mb_x > 0) |
235 | ADD_PRED(preds[0], s->mv_table[0][mb_i - s->b_width - 1][dir][0], s->mv_table[0][mb_i - s->b_width - 1][dir][1]); |
236 | } |
237 | |
238 | //median predictor |
239 | if (preds[0].nb == 4) { |
240 | me_ctx->pred_x = mid_pred(preds[0].mvs[1][0], preds[0].mvs[2][0], preds[0].mvs[3][0]); |
241 | me_ctx->pred_y = mid_pred(preds[0].mvs[1][1], preds[0].mvs[2][1], preds[0].mvs[3][1]); |
242 | } else if (preds[0].nb == 3) { |
243 | me_ctx->pred_x = mid_pred(0, preds[0].mvs[1][0], preds[0].mvs[2][0]); |
244 | me_ctx->pred_y = mid_pred(0, preds[0].mvs[1][1], preds[0].mvs[2][1]); |
245 | } else if (preds[0].nb == 2) { |
246 | me_ctx->pred_x = preds[0].mvs[1][0]; |
247 | me_ctx->pred_y = preds[0].mvs[1][1]; |
248 | } else { |
249 | me_ctx->pred_x = 0; |
250 | me_ctx->pred_y = 0; |
251 | } |
252 | |
253 | ff_me_search_umh(me_ctx, x_mb, y_mb, mv); |
254 | |
255 | s->mv_table[0][mb_i][dir][0] = mv[0] - x_mb; |
256 | s->mv_table[0][mb_i][dir][1] = mv[1] - y_mb; |
257 | add_mv_data(((AVMotionVector *) sd->data) + mv_count++, me_ctx->mb_size, x_mb, y_mb, mv[0], mv[1], dir); |
258 | } |
259 | |
260 | } else if (s->method == AV_ME_METHOD_EPZS) { |
261 | |
262 | for (mb_y = 0; mb_y < s->b_height; mb_y++) |
263 | for (mb_x = 0; mb_x < s->b_width; mb_x++) { |
264 | const int mb_i = mb_x + mb_y * s->b_width; |
265 | const int x_mb = mb_x << s->log2_mb_size; |
266 | const int y_mb = mb_y << s->log2_mb_size; |
267 | int mv[2] = {x_mb, y_mb}; |
268 | |
269 | AVMotionEstPredictor *preds = me_ctx->preds; |
270 | preds[0].nb = 0; |
271 | preds[1].nb = 0; |
272 | |
273 | ADD_PRED(preds[0], 0, 0); |
274 | |
275 | //left mb in current frame |
276 | if (mb_x > 0) |
277 | ADD_PRED(preds[0], s->mv_table[0][mb_i - 1][dir][0], s->mv_table[0][mb_i - 1][dir][1]); |
278 | |
279 | //top mb in current frame |
280 | if (mb_y > 0) |
281 | ADD_PRED(preds[0], s->mv_table[0][mb_i - s->b_width][dir][0], s->mv_table[0][mb_i - s->b_width][dir][1]); |
282 | |
283 | //top-right mb in current frame |
284 | if (mb_y > 0 && mb_x + 1 < s->b_width) |
285 | ADD_PRED(preds[0], s->mv_table[0][mb_i - s->b_width + 1][dir][0], s->mv_table[0][mb_i - s->b_width + 1][dir][1]); |
286 | |
287 | //median predictor |
288 | if (preds[0].nb == 4) { |
289 | me_ctx->pred_x = mid_pred(preds[0].mvs[1][0], preds[0].mvs[2][0], preds[0].mvs[3][0]); |
290 | me_ctx->pred_y = mid_pred(preds[0].mvs[1][1], preds[0].mvs[2][1], preds[0].mvs[3][1]); |
291 | } else if (preds[0].nb == 3) { |
292 | me_ctx->pred_x = mid_pred(0, preds[0].mvs[1][0], preds[0].mvs[2][0]); |
293 | me_ctx->pred_y = mid_pred(0, preds[0].mvs[1][1], preds[0].mvs[2][1]); |
294 | } else if (preds[0].nb == 2) { |
295 | me_ctx->pred_x = preds[0].mvs[1][0]; |
296 | me_ctx->pred_y = preds[0].mvs[1][1]; |
297 | } else { |
298 | me_ctx->pred_x = 0; |
299 | me_ctx->pred_y = 0; |
300 | } |
301 | |
302 | //collocated mb in prev frame |
303 | ADD_PRED(preds[0], s->mv_table[1][mb_i][dir][0], s->mv_table[1][mb_i][dir][1]); |
304 | |
305 | //accelerator motion vector of collocated block in prev frame |
306 | ADD_PRED(preds[1], s->mv_table[1][mb_i][dir][0] + (s->mv_table[1][mb_i][dir][0] - s->mv_table[2][mb_i][dir][0]), |
307 | s->mv_table[1][mb_i][dir][1] + (s->mv_table[1][mb_i][dir][1] - s->mv_table[2][mb_i][dir][1])); |
308 | |
309 | //left mb in prev frame |
310 | if (mb_x > 0) |
311 | ADD_PRED(preds[1], s->mv_table[1][mb_i - 1][dir][0], s->mv_table[1][mb_i - 1][dir][1]); |
312 | |
313 | //top mb in prev frame |
314 | if (mb_y > 0) |
315 | ADD_PRED(preds[1], s->mv_table[1][mb_i - s->b_width][dir][0], s->mv_table[1][mb_i - s->b_width][dir][1]); |
316 | |
317 | //right mb in prev frame |
318 | if (mb_x + 1 < s->b_width) |
319 | ADD_PRED(preds[1], s->mv_table[1][mb_i + 1][dir][0], s->mv_table[1][mb_i + 1][dir][1]); |
320 | |
321 | //bottom mb in prev frame |
322 | if (mb_y + 1 < s->b_height) |
323 | ADD_PRED(preds[1], s->mv_table[1][mb_i + s->b_width][dir][0], s->mv_table[1][mb_i + s->b_width][dir][1]); |
324 | |
325 | ff_me_search_epzs(me_ctx, x_mb, y_mb, mv); |
326 | |
327 | s->mv_table[0][mb_i][dir][0] = mv[0] - x_mb; |
328 | s->mv_table[0][mb_i][dir][1] = mv[1] - y_mb; |
329 | add_mv_data(((AVMotionVector *) sd->data) + mv_count++, s->mb_size, x_mb, y_mb, mv[0], mv[1], dir); |
330 | } |
331 | } |
332 | } |
333 | |
334 | return ff_filter_frame(ctx->outputs[0], out); |
335 | } |
336 | |
337 | static av_cold void uninit(AVFilterContext *ctx) |
338 | { |
339 | MEContext *s = ctx->priv; |
340 | int i; |
341 | |
342 | av_frame_free(&s->prev); |
343 | av_frame_free(&s->cur); |
344 | av_frame_free(&s->next); |
345 | |
346 | for (i = 0; i < 3; i++) |
347 | av_freep(&s->mv_table[i]); |
348 | } |
349 | |
350 | static const AVFilterPad mestimate_inputs[] = { |
351 | { |
352 | .name = "default", |
353 | .type = AVMEDIA_TYPE_VIDEO, |
354 | .filter_frame = filter_frame, |
355 | .config_props = config_input, |
356 | }, |
357 | { NULL } |
358 | }; |
359 | |
360 | static const AVFilterPad mestimate_outputs[] = { |
361 | { |
362 | .name = "default", |
363 | .type = AVMEDIA_TYPE_VIDEO, |
364 | }, |
365 | { NULL } |
366 | }; |
367 | |
368 | AVFilter ff_vf_mestimate = { |
369 | .name = "mestimate", |
370 | .description = NULL_IF_CONFIG_SMALL("Generate motion vectors."), |
371 | .priv_size = sizeof(MEContext), |
372 | .priv_class = &mestimate_class, |
373 | .uninit = uninit, |
374 | .query_formats = query_formats, |
375 | .inputs = mestimate_inputs, |
376 | .outputs = mestimate_outputs, |
377 | }; |
378 |