lavfi: add select filter
[ffmpeg.git] / libavfilter / vf_select.c
1 /*
2  * Copyright (c) 2011 Stefano Sabatini
3  *
4  * This file is part of Libav.
5  *
6  * Libav 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  * Libav 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 Libav; 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  * filter for selecting which frame passes in the filterchain
24  */
25
26 #include "libavutil/eval.h"
27 #include "libavutil/fifo.h"
28 #include "libavutil/mathematics.h"
29 #include "avfilter.h"
30
31 static const char *var_names[] = {
32     "E",                 ///< Euler number
33     "PHI",               ///< golden ratio
34     "PI",                ///< greek pi
35
36     "TB",                ///< timebase
37
38     "pts",               ///< original pts in the file of the frame
39     "start_pts",         ///< first PTS in the stream, expressed in TB units
40     "prev_pts",          ///< previous frame PTS
41     "prev_selected_pts", ///< previous selected frame PTS
42
43     "t",                 ///< first PTS in seconds
44     "start_t",           ///< first PTS in the stream, expressed in seconds
45     "prev_t",            ///< previous frame time
46     "prev_selected_t",   ///< previously selected time
47
48     "pict_type",         ///< the type of picture in the movie
49     "I",
50     "P",
51     "B",
52     "S",
53     "SI",
54     "SP",
55     "BI",
56
57     "interlace_type",    ///< the frame interlace type
58     "PROGRESSIVE",
59     "TOPFIRST",
60     "BOTTOMFIRST",
61
62     "n",                 ///< frame number (starting from zero)
63     "selected_n",        ///< selected frame number (starting from zero)
64     "prev_selected_n",   ///< number of the last selected frame
65
66     "key",               ///< tell if the frame is a key frame
67     "pos",               ///< original position in the file of the frame
68
69     NULL
70 };
71
72 enum var_name {
73     VAR_E,
74     VAR_PHI,
75     VAR_PI,
76
77     VAR_TB,
78
79     VAR_PTS,
80     VAR_START_PTS,
81     VAR_PREV_PTS,
82     VAR_PREV_SELECTED_PTS,
83
84     VAR_T,
85     VAR_START_T,
86     VAR_PREV_T,
87     VAR_PREV_SELECTED_T,
88
89     VAR_PICT_TYPE,
90     VAR_PICT_TYPE_I,
91     VAR_PICT_TYPE_P,
92     VAR_PICT_TYPE_B,
93     VAR_PICT_TYPE_S,
94     VAR_PICT_TYPE_SI,
95     VAR_PICT_TYPE_SP,
96     VAR_PICT_TYPE_BI,
97
98     VAR_INTERLACE_TYPE,
99     VAR_INTERLACE_TYPE_P,
100     VAR_INTERLACE_TYPE_T,
101     VAR_INTERLACE_TYPE_B,
102
103     VAR_N,
104     VAR_SELECTED_N,
105     VAR_PREV_SELECTED_N,
106
107     VAR_KEY,
108     VAR_POS,
109
110     VAR_VARS_NB
111 };
112
113 #define FIFO_SIZE 8
114
115 typedef struct {
116     AVExpr *expr;
117     double var_values[VAR_VARS_NB];
118     double select;
119     int cache_frames;
120     AVFifoBuffer *pending_frames; ///< FIFO buffer of video frames
121 } SelectContext;
122
123 static av_cold int init(AVFilterContext *ctx, const char *args, void *opaque)
124 {
125     SelectContext *select = ctx->priv;
126     int ret;
127
128     if ((ret = av_expr_parse(&select->expr, args ? args : "1",
129                              var_names, NULL, NULL, NULL, NULL, 0, ctx)) < 0) {
130         av_log(ctx, AV_LOG_ERROR, "Error while parsing expression '%s'\n", args);
131         return ret;
132     }
133
134     select->pending_frames = av_fifo_alloc(FIFO_SIZE*sizeof(AVFilterBufferRef*));
135     if (!select->pending_frames) {
136         av_log(ctx, AV_LOG_ERROR, "Failed to allocate pending frames buffer.\n");
137         return AVERROR(ENOMEM);
138     }
139     return 0;
140 }
141
142 #define INTERLACE_TYPE_P 0
143 #define INTERLACE_TYPE_T 1
144 #define INTERLACE_TYPE_B 2
145
146 static int config_input(AVFilterLink *inlink)
147 {
148     SelectContext *select = inlink->dst->priv;
149
150     select->var_values[VAR_E]   = M_E;
151     select->var_values[VAR_PHI] = M_PHI;
152     select->var_values[VAR_PI]  = M_PI;
153
154     select->var_values[VAR_N]          = 0.0;
155     select->var_values[VAR_SELECTED_N] = 0.0;
156
157     select->var_values[VAR_TB] = av_q2d(inlink->time_base);
158
159     select->var_values[VAR_PREV_PTS]          = NAN;
160     select->var_values[VAR_PREV_SELECTED_PTS] = NAN;
161     select->var_values[VAR_PREV_SELECTED_T]   = NAN;
162     select->var_values[VAR_START_PTS]         = NAN;
163     select->var_values[VAR_START_T]           = NAN;
164
165     select->var_values[VAR_PICT_TYPE_I]  = AV_PICTURE_TYPE_I;
166     select->var_values[VAR_PICT_TYPE_P]  = AV_PICTURE_TYPE_P;
167     select->var_values[VAR_PICT_TYPE_B]  = AV_PICTURE_TYPE_B;
168     select->var_values[VAR_PICT_TYPE_SI] = AV_PICTURE_TYPE_SI;
169     select->var_values[VAR_PICT_TYPE_SP] = AV_PICTURE_TYPE_SP;
170
171     select->var_values[VAR_INTERLACE_TYPE_P] = INTERLACE_TYPE_P;
172     select->var_values[VAR_INTERLACE_TYPE_T] = INTERLACE_TYPE_T;
173     select->var_values[VAR_INTERLACE_TYPE_B] = INTERLACE_TYPE_B;;
174
175     return 0;
176 }
177
178 #define D2TS(d)  (isnan(d) ? AV_NOPTS_VALUE : (int64_t)(d))
179 #define TS2D(ts) ((ts) == AV_NOPTS_VALUE ? NAN : (double)(ts))
180
181 static int select_frame(AVFilterContext *ctx, AVFilterBufferRef *picref)
182 {
183     SelectContext *select = ctx->priv;
184     AVFilterLink *inlink = ctx->inputs[0];
185     double res;
186
187     if (isnan(select->var_values[VAR_START_PTS]))
188         select->var_values[VAR_START_PTS] = TS2D(picref->pts);
189     if (isnan(select->var_values[VAR_START_T]))
190         select->var_values[VAR_START_T] = TS2D(picref->pts) * av_q2d(inlink->time_base);
191
192     select->var_values[VAR_PTS] = TS2D(picref->pts);
193     select->var_values[VAR_T  ] = TS2D(picref->pts) * av_q2d(inlink->time_base);
194     select->var_values[VAR_POS] = picref->pos == -1 ? NAN : picref->pos;
195     select->var_values[VAR_PREV_PTS] = TS2D(picref ->pts);
196
197     select->var_values[VAR_INTERLACE_TYPE] =
198         !picref->video->interlaced     ? INTERLACE_TYPE_P :
199         picref->video->top_field_first ? INTERLACE_TYPE_T : INTERLACE_TYPE_B;
200     select->var_values[VAR_PICT_TYPE] = picref->video->pict_type;
201
202     res = av_expr_eval(select->expr, select->var_values, NULL);
203     av_log(inlink->dst, AV_LOG_DEBUG,
204            "n:%d pts:%d t:%f pos:%d interlace_type:%c key:%d pict_type:%c "
205            "-> select:%f\n",
206            (int)select->var_values[VAR_N],
207            (int)select->var_values[VAR_PTS],
208            select->var_values[VAR_T],
209            (int)select->var_values[VAR_POS],
210            select->var_values[VAR_INTERLACE_TYPE] == INTERLACE_TYPE_P ? 'P' :
211            select->var_values[VAR_INTERLACE_TYPE] == INTERLACE_TYPE_T ? 'T' :
212            select->var_values[VAR_INTERLACE_TYPE] == INTERLACE_TYPE_B ? 'B' : '?',
213            (int)select->var_values[VAR_KEY],
214            av_get_picture_type_char(select->var_values[VAR_PICT_TYPE]),
215            res);
216
217     select->var_values[VAR_N] += 1.0;
218
219     if (res) {
220         select->var_values[VAR_PREV_SELECTED_N]   = select->var_values[VAR_N];
221         select->var_values[VAR_PREV_SELECTED_PTS] = select->var_values[VAR_PTS];
222         select->var_values[VAR_PREV_SELECTED_T]   = select->var_values[VAR_T];
223         select->var_values[VAR_SELECTED_N] += 1.0;
224     }
225     return res;
226 }
227
228 static void start_frame(AVFilterLink *inlink, AVFilterBufferRef *picref)
229 {
230     SelectContext *select = inlink->dst->priv;
231
232     select->select = select_frame(inlink->dst, picref);
233     if (select->select) {
234         /* frame was requested through poll_frame */
235         if (select->cache_frames) {
236             if (!av_fifo_space(select->pending_frames))
237                 av_log(inlink->dst, AV_LOG_ERROR,
238                        "Buffering limit reached, cannot cache more frames\n");
239             else
240                 av_fifo_generic_write(select->pending_frames, &picref,
241                                       sizeof(picref), NULL);
242             return;
243         }
244         avfilter_start_frame(inlink->dst->outputs[0], avfilter_ref_buffer(picref, ~0));
245     }
246 }
247
248 static void draw_slice(AVFilterLink *inlink, int y, int h, int slice_dir)
249 {
250     SelectContext *select = inlink->dst->priv;
251
252     if (select->select && !select->cache_frames)
253         avfilter_draw_slice(inlink->dst->outputs[0], y, h, slice_dir);
254 }
255
256 static void end_frame(AVFilterLink *inlink)
257 {
258     SelectContext *select = inlink->dst->priv;
259     AVFilterBufferRef *picref = inlink->cur_buf;
260
261     if (select->select) {
262         if (select->cache_frames)
263             return;
264         avfilter_end_frame(inlink->dst->outputs[0]);
265     }
266     avfilter_unref_buffer(picref);
267 }
268
269 static int request_frame(AVFilterLink *outlink)
270 {
271     AVFilterContext *ctx = outlink->src;
272     SelectContext *select = ctx->priv;
273     AVFilterLink *inlink = outlink->src->inputs[0];
274     select->select = 0;
275
276     if (av_fifo_size(select->pending_frames)) {
277         AVFilterBufferRef *picref;
278         av_fifo_generic_read(select->pending_frames, &picref, sizeof(picref), NULL);
279         avfilter_start_frame(outlink, avfilter_ref_buffer(picref, ~0));
280         avfilter_draw_slice(outlink, 0, outlink->h, 1);
281         avfilter_end_frame(outlink);
282         avfilter_unref_buffer(picref);
283         return 0;
284     }
285
286     while (!select->select) {
287         int ret = avfilter_request_frame(inlink);
288         if (ret < 0)
289             return ret;
290     }
291
292     return 0;
293 }
294
295 static int poll_frame(AVFilterLink *outlink)
296 {
297     SelectContext *select = outlink->src->priv;
298     AVFilterLink *inlink = outlink->src->inputs[0];
299     int count, ret;
300
301     if (!av_fifo_size(select->pending_frames)) {
302         if ((count = avfilter_poll_frame(inlink)) <= 0)
303             return count;
304         /* request frame from input, and apply select condition to it */
305         select->cache_frames = 1;
306         while (count-- && av_fifo_space(select->pending_frames)) {
307             ret = avfilter_request_frame(inlink);
308             if (ret < 0)
309                 break;
310         }
311         select->cache_frames = 0;
312     }
313
314     return av_fifo_size(select->pending_frames)/sizeof(AVFilterBufferRef *);
315 }
316
317 static av_cold void uninit(AVFilterContext *ctx)
318 {
319     SelectContext *select = ctx->priv;
320     AVFilterBufferRef *picref;
321
322     av_expr_free(select->expr);
323     select->expr = NULL;
324
325     while (select->pending_frames &&
326            av_fifo_generic_read(select->pending_frames, &picref, sizeof(picref), NULL) == sizeof(picref))
327         avfilter_unref_buffer(picref);
328     av_fifo_free(select->pending_frames);
329     select->pending_frames = NULL;
330 }
331
332 AVFilter avfilter_vf_select = {
333     .name      = "select",
334     .description = NULL_IF_CONFIG_SMALL("Select frames to pass in output."),
335     .init      = init,
336     .uninit    = uninit,
337
338     .priv_size = sizeof(SelectContext),
339
340     .inputs    = (AVFilterPad[]) {{ .name             = "default",
341                                     .type             = AVMEDIA_TYPE_VIDEO,
342                                     .get_video_buffer = avfilter_null_get_video_buffer,
343                                     .config_props     = config_input,
344                                     .start_frame      = start_frame,
345                                     .draw_slice       = draw_slice,
346                                     .end_frame        = end_frame },
347                                   { .name = NULL }},
348     .outputs   = (AVFilterPad[]) {{ .name             = "default",
349                                     .type             = AVMEDIA_TYPE_VIDEO,
350                                     .poll_frame       = poll_frame,
351                                     .request_frame    = request_frame, },
352                                   { .name = NULL}},
353 };