1 /*
2  * Copyright (c) 2021 Paul B Mahol
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 /**
22  * @file
23  * Change the PTS/DTS timestamps.
24  */
25 
26 #include "libavutil/opt.h"
27 #include "libavutil/eval.h"
28 
29 #include "bsf.h"
30 #include "bsf_internal.h"
31 
32 static const char *const var_names[] = {
33     "N",           ///< frame number (starting at zero)
34     "TS",
35     "POS",         ///< original position in the file of the frame
36     "PREV_INPTS",  ///< previous  input PTS
37     "PREV_INDTS",  ///< previous  input DTS
38     "PREV_INDURATION", ///< previous input duration
39     "PREV_OUTPTS", ///< previous output PTS
40     "PREV_OUTDTS", ///< previous output DTS
41     "PREV_OUTDURATION", ///< previous output duration
42     "NEXT_PTS",    ///< next input PTS
43     "NEXT_DTS",    ///< next input DTS
44     "NEXT_DURATION", ///< next input duration
45     "PTS",         ///< original PTS in the file of the frame
46     "DTS",         ///< original DTS in the file of the frame
47     "DURATION",    ///< original duration in the file of the frame
48     "STARTPTS",    ///< PTS at start of movie
49     "STARTDTS",    ///< DTS at start of movie
50     "TB",          ///< input timebase of the stream
51     "TB_OUT",      ///< output timebase of the stream
52     "SR",          ///< sample rate of the stream
53     "NOPTS",       ///< The AV_NOPTS_VALUE constant
54     NULL
55 };
56 
57 enum var_name {
58     VAR_N,
59     VAR_TS,
60     VAR_POS,
61     VAR_PREV_INPTS,
62     VAR_PREV_INDTS,
63     VAR_PREV_INDUR,
64     VAR_PREV_OUTPTS,
65     VAR_PREV_OUTDTS,
66     VAR_PREV_OUTDUR,
67     VAR_NEXT_PTS,
68     VAR_NEXT_DTS,
69     VAR_NEXT_DUR,
70     VAR_PTS,
71     VAR_DTS,
72     VAR_DURATION,
73     VAR_STARTPTS,
74     VAR_STARTDTS,
75     VAR_TB,
76     VAR_TB_OUT,
77     VAR_SR,
78     VAR_NOPTS,
79     VAR_VARS_NB
80 };
81 
82 typedef struct SetTSContext {
83     const AVClass *class;
84 
85     char *ts_str;
86     char *pts_str;
87     char *dts_str;
88     char *duration_str;
89 
90     AVRational time_base;
91 
92     int64_t frame_number;
93 
94     double var_values[VAR_VARS_NB];
95 
96     AVExpr *ts_expr;
97     AVExpr *pts_expr;
98     AVExpr *dts_expr;
99     AVExpr *duration_expr;
100 
101     AVPacket *prev_inpkt;
102     AVPacket *prev_outpkt;
103     AVPacket *cur_pkt;
104 } SetTSContext;
105 
setts_init(AVBSFContext *ctx)106 static int setts_init(AVBSFContext *ctx)
107 {
108     SetTSContext *s = ctx->priv_data;
109     int ret;
110 
111     s->prev_inpkt = av_packet_alloc();
112     s->prev_outpkt = av_packet_alloc();
113     s->cur_pkt = av_packet_alloc();
114     if (!s->prev_inpkt || !s->prev_outpkt || !s->cur_pkt)
115         return AVERROR(ENOMEM);
116 
117     if ((ret = av_expr_parse(&s->ts_expr, s->ts_str,
118                              var_names, NULL, NULL, NULL, NULL, 0, ctx)) < 0) {
119         av_log(ctx, AV_LOG_ERROR, "Error while parsing ts expression '%s'\n", s->ts_str);
120         return ret;
121     }
122 
123     if ((ret = av_expr_parse(&s->duration_expr, s->duration_str,
124                              var_names, NULL, NULL, NULL, NULL, 0, ctx)) < 0) {
125         av_log(ctx, AV_LOG_ERROR, "Error while parsing duration expression '%s'\n", s->duration_str);
126         return ret;
127     }
128 
129     if (s->pts_str) {
130         if ((ret = av_expr_parse(&s->pts_expr, s->pts_str,
131                                  var_names, NULL, NULL, NULL, NULL, 0, ctx)) < 0) {
132             av_log(ctx, AV_LOG_ERROR, "Error while parsing pts expression '%s'\n", s->pts_str);
133             return ret;
134         }
135     }
136 
137     if (s->dts_str) {
138         if ((ret = av_expr_parse(&s->dts_expr, s->dts_str,
139                                  var_names, NULL, NULL, NULL, NULL, 0, ctx)) < 0) {
140             av_log(ctx, AV_LOG_ERROR, "Error while parsing dts expression '%s'\n", s->dts_str);
141             return ret;
142         }
143     }
144 
145     if (s->time_base.num > 0 && s->time_base.den > 0)
146         ctx->time_base_out = s->time_base;
147 
148     s->frame_number= 0;
149     s->var_values[VAR_STARTPTS] = AV_NOPTS_VALUE;
150     s->var_values[VAR_STARTDTS] = AV_NOPTS_VALUE;
151     s->var_values[VAR_NOPTS] = AV_NOPTS_VALUE;
152     s->var_values[VAR_TB]    = ctx->time_base_in.den ? av_q2d(ctx->time_base_in) : 0;
153     s->var_values[VAR_TB_OUT]= ctx->time_base_out.den ? av_q2d(ctx->time_base_out) : 0;
154     s->var_values[VAR_SR]    = ctx->par_in->sample_rate;
155 
156     return 0;
157 }
158 
setts_filter(AVBSFContext *ctx, AVPacket *pkt)159 static int setts_filter(AVBSFContext *ctx, AVPacket *pkt)
160 {
161     SetTSContext *s = ctx->priv_data;
162     int64_t new_ts, new_pts, new_dts, new_duration;
163     int ret;
164 
165     ret = ff_bsf_get_packet_ref(ctx, pkt);
166     if (ret < 0 && (ret != AVERROR_EOF || !s->cur_pkt->data))
167         return ret;
168 
169     if (!s->cur_pkt->data) {
170          av_packet_move_ref(s->cur_pkt, pkt);
171          return AVERROR(EAGAIN);
172     }
173 
174     if (s->var_values[VAR_STARTPTS] == AV_NOPTS_VALUE)
175         s->var_values[VAR_STARTPTS] = s->cur_pkt->pts;
176 
177     if (s->var_values[VAR_STARTDTS] == AV_NOPTS_VALUE)
178         s->var_values[VAR_STARTDTS] = s->cur_pkt->dts;
179 
180     s->var_values[VAR_N]           = s->frame_number++;
181     s->var_values[VAR_TS]          = s->cur_pkt->dts;
182     s->var_values[VAR_POS]         = s->cur_pkt->pos;
183     s->var_values[VAR_PTS]         = s->cur_pkt->pts;
184     s->var_values[VAR_DTS]         = s->cur_pkt->dts;
185     s->var_values[VAR_DURATION]    = s->cur_pkt->duration;
186     s->var_values[VAR_PREV_INPTS]  = s->prev_inpkt->pts;
187     s->var_values[VAR_PREV_INDTS]  = s->prev_inpkt->dts;
188     s->var_values[VAR_PREV_INDUR]  = s->prev_inpkt->duration;
189     s->var_values[VAR_PREV_OUTPTS] = s->prev_outpkt->pts;
190     s->var_values[VAR_PREV_OUTDTS] = s->prev_outpkt->dts;
191     s->var_values[VAR_PREV_OUTDUR] = s->prev_outpkt->duration;
192     s->var_values[VAR_NEXT_PTS]    = pkt->pts;
193     s->var_values[VAR_NEXT_DTS]    = pkt->dts;
194     s->var_values[VAR_NEXT_DUR]    = pkt->duration;
195 
196     new_ts = llrint(av_expr_eval(s->ts_expr, s->var_values, NULL));
197     new_duration = llrint(av_expr_eval(s->duration_expr, s->var_values, NULL));
198 
199     if (s->pts_str) {
200         s->var_values[VAR_TS] = s->cur_pkt->pts;
201         new_pts = llrint(av_expr_eval(s->pts_expr, s->var_values, NULL));
202     } else {
203         new_pts = new_ts;
204     }
205 
206     if (s->dts_str) {
207         s->var_values[VAR_TS] = s->cur_pkt->dts;
208         new_dts = llrint(av_expr_eval(s->dts_expr, s->var_values, NULL));
209     } else {
210         new_dts = new_ts;
211     }
212 
213     av_packet_unref(s->prev_inpkt);
214     av_packet_unref(s->prev_outpkt);
215     av_packet_move_ref(s->prev_inpkt, s->cur_pkt);
216     av_packet_move_ref(s->cur_pkt, pkt);
217 
218     ret = av_packet_ref(pkt, s->prev_inpkt);
219     if (ret < 0)
220         return ret;
221 
222     pkt->pts = new_pts;
223     pkt->dts = new_dts;
224     pkt->duration = new_duration;
225 
226     ret = av_packet_ref(s->prev_outpkt, pkt);
227     if (ret < 0)
228         av_packet_unref(pkt);
229 
230     return ret;
231 }
232 
setts_close(AVBSFContext *bsf)233 static void setts_close(AVBSFContext *bsf)
234 {
235     SetTSContext *s = bsf->priv_data;
236 
237     av_packet_free(&s->prev_inpkt);
238     av_packet_free(&s->prev_outpkt);
239     av_packet_free(&s->cur_pkt);
240 
241     av_expr_free(s->ts_expr);
242     s->ts_expr = NULL;
243     av_expr_free(s->pts_expr);
244     s->pts_expr = NULL;
245     av_expr_free(s->dts_expr);
246     s->dts_expr = NULL;
247     av_expr_free(s->duration_expr);
248     s->duration_expr = NULL;
249 }
250 
251 #define OFFSET(x) offsetof(SetTSContext, x)
252 #define FLAGS (AV_OPT_FLAG_VIDEO_PARAM|AV_OPT_FLAG_AUDIO_PARAM|AV_OPT_FLAG_SUBTITLE_PARAM|AV_OPT_FLAG_BSF_PARAM)
253 
254 static const AVOption options[] = {
255     { "ts",  "set expression for packet PTS and DTS", OFFSET(ts_str),  AV_OPT_TYPE_STRING, {.str="TS"}, 0, 0, FLAGS },
256     { "pts", "set expression for packet PTS", OFFSET(pts_str), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS },
257     { "dts", "set expression for packet DTS", OFFSET(dts_str), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS },
258     { "duration", "set expression for packet duration", OFFSET(duration_str), AV_OPT_TYPE_STRING, {.str="DURATION"}, 0, 0, FLAGS },
259     { "time_base", "set output timebase", OFFSET(time_base), AV_OPT_TYPE_RATIONAL, {.dbl=0}, 0, INT_MAX, FLAGS },
260     { NULL },
261 };
262 
263 static const AVClass setts_class = {
264     .class_name = "setts_bsf",
265     .item_name  = av_default_item_name,
266     .option     = options,
267     .version    = LIBAVUTIL_VERSION_INT,
268 };
269 
270 const FFBitStreamFilter ff_setts_bsf = {
271     .p.name         = "setts",
272     .p.priv_class   = &setts_class,
273     .priv_data_size = sizeof(SetTSContext),
274     .init           = setts_init,
275     .close          = setts_close,
276     .filter         = setts_filter,
277 };
278