ffmpeg: implement -force_key_frames expression evalution

3 files changed, 112 insertions, 5 deletions

diff --git a/doc/ffmpeg.texi b/doc/ffmpeg.texi
index 2c023e1..034ca13 100644
--- a/doc/ffmpeg.texi
+++ b/doc/ffmpeg.texi
@@ -552,9 +552,16 @@ Force video tag/fourcc. This is an alias for @code{-tag:v}.
 Show QP histogram
 @item -vbsf @var{bitstream_filter}
 Deprecated see -bsf
+
 @item -force_key_frames[:@var{stream_specifier}] @var{time}[,@var{time}...] (@emph{output,per-stream})
+@item -force_key_frames[:@var{stream_specifier}] expr:@var{expr} (@emph{output,per-stream})
 Force key frames at the specified timestamps, more precisely at the first
 frames after each specified time.
+
+If the argument is prefixed with @code{expr:}, the string @var{expr}
+is interpreted like an expression and is evaluated for each frame. A
+key frame is forced in case the evaluation is non-zero.
+
 If one of the times is "@code{chapters}[@var{delta}]", it is expanded into
 the time of the beginning of all chapters in the file, shifted by
 @var{delta}, expressed as a time in seconds.
@@ -567,6 +574,37 @@ before the beginning of every chapter:
 -force_key_frames 0:05:00,chapters-0.1
 @end example
 
+The expression in @var{expr} can contain the following constants:
+@table @option
+@item n
+the number of current processed frame, starting from 0
+@item n_forced
+the number of forced frames
+@item prev_forced_n
+the number of the previous forced frame, it is @code{NAN} when no
+keyframe was forced yet
+@item prev_forced_t
+the time of the previous forced frame, it is @code{NAN} when no
+keyframe was forced yet
+@item t
+the time of the current processed frame
+@end table
+
+For example to force a key frame every 5 seconds, you can specify:
+@example
+-force_key_frames expr:gte(t,n_forced*5)
+@end example
+
+To force a key frame 5 seconds after the time of the last forced one,
+starting from second 13:
+@example
+-force_key_frames expr:if(isnan(prev_forced_t),gte(t,13),gte(t,prev_forced_t+5))
+@end example
+
+Note that forcing too many keyframes is very harmful for the lookahead
+algorithms of certain encoders: using fixed-GOP options or similar
+would be more efficient.
+
 @item -copyinkf[:@var{stream_specifier}] (@emph{output,per-stream})
 When doing stream copy, copy also non-key frames found at the
 beginning.
diff --git a/ffmpeg.c b/ffmpeg.c
index 54097f4..bbd21b6 100644
--- a/ffmpeg.c
+++ b/ffmpeg.c
@@ -109,6 +109,15 @@ const int program_birth_year = 2000;
 
 static FILE *vstats_file;
 
+const char *const forced_keyframes_const_names[] = {
+    "n",
+    "n_forced",
+    "prev_forced_n",
+    "prev_forced_t",
+    "t",
+    NULL
+};
+
 static void do_video_stats(OutputStream *ost, int frame_size);
 static int64_t getutime(void);
 
@@ -437,6 +446,7 @@ static void exit_program(void)
         avcodec_free_frame(&output_streams[i]->filtered_frame);
 
         av_freep(&output_streams[i]->forced_keyframes);
+        av_expr_free(output_streams[i]->forced_keyframes_pexpr);
         av_freep(&output_streams[i]->avfilter);
         av_freep(&output_streams[i]->logfile_prefix);
         av_freep(&output_streams[i]);
@@ -873,8 +883,9 @@ static void do_video_out(AVFormatContext *s,
         video_size += pkt.size;
         write_frame(s, &pkt, ost);
     } else {
-        int got_packet;
+        int got_packet, forced_keyframe = 0;
         AVFrame big_picture;
+        double pts_time;
 
         big_picture = *in_picture;
         /* better than nothing: use input picture interlaced
@@ -898,11 +909,41 @@ static void do_video_out(AVFormatContext *s,
         big_picture.quality = ost->st->codec->global_quality;
         if (!enc->me_threshold)
             big_picture.pict_type = 0;
+
+        pts_time = big_picture.pts != AV_NOPTS_VALUE ?
+            big_picture.pts * av_q2d(enc->time_base) : NAN;
         if (ost->forced_kf_index < ost->forced_kf_count &&
             big_picture.pts >= ost->forced_kf_pts[ost->forced_kf_index]) {
-            big_picture.pict_type = AV_PICTURE_TYPE_I;
             ost->forced_kf_index++;
+            forced_keyframe = 1;
+        } else if (ost->forced_keyframes_pexpr) {
+            double res;
+            ost->forced_keyframes_expr_const_values[FKF_T] = pts_time;
+            res = av_expr_eval(ost->forced_keyframes_pexpr,
+                               ost->forced_keyframes_expr_const_values, NULL);
+            av_dlog(NULL, "force_key_frame: n:%f n_forced:%f prev_forced_n:%f t:%f prev_forced_t:%f -> res:%f\n",
+                    ost->forced_keyframes_expr_const_values[FKF_N],
+                    ost->forced_keyframes_expr_const_values[FKF_N_FORCED],
+                    ost->forced_keyframes_expr_const_values[FKF_PREV_FORCED_N],
+                    ost->forced_keyframes_expr_const_values[FKF_T],
+                    ost->forced_keyframes_expr_const_values[FKF_PREV_FORCED_T],
+                    res);
+            if (res) {
+                forced_keyframe = 1;
+                ost->forced_keyframes_expr_const_values[FKF_PREV_FORCED_N] =
+                    ost->forced_keyframes_expr_const_values[FKF_N];
+                ost->forced_keyframes_expr_const_values[FKF_PREV_FORCED_T] =
+                    ost->forced_keyframes_expr_const_values[FKF_T];
+                ost->forced_keyframes_expr_const_values[FKF_N_FORCED] += 1;
+            }
+
+            ost->forced_keyframes_expr_const_values[FKF_N] += 1;
+        }
+        if (forced_keyframe) {
+            big_picture.pict_type = AV_PICTURE_TYPE_I;
+            av_log(NULL, AV_LOG_DEBUG, "Forced keyframe at time %f\n", pts_time);
         }
+
         update_benchmark(NULL);
         ret = avcodec_encode_video2(enc, &pkt, &big_picture, &got_packet);
         update_benchmark("encode_video %d.%d", ost->file_index, ost->index);
@@ -2272,9 +2313,23 @@ static int transcode_init(void)
                     codec->bits_per_raw_sample = frame_bits_per_raw_sample;
                 }
 
-                if (ost->forced_keyframes)
-                    parse_forced_key_frames(ost->forced_keyframes, ost,
-                                            ost->st->codec);
+                if (ost->forced_keyframes) {
+                    if (!strncmp(ost->forced_keyframes, "expr:", 5)) {
+                        ret = av_expr_parse(&ost->forced_keyframes_pexpr, ost->forced_keyframes+5,
+                                            forced_keyframes_const_names, NULL, NULL, NULL, NULL, 0, NULL);
+                        if (ret < 0) {
+                            av_log(NULL, AV_LOG_ERROR,
+                                   "Invalid force_key_frames expression '%s'\n", ost->forced_keyframes+5);
+                            return ret;
+                        }
+                        ost->forced_keyframes_expr_const_values[FKF_N] = 0;
+                        ost->forced_keyframes_expr_const_values[FKF_N_FORCED] = 0;
+                        ost->forced_keyframes_expr_const_values[FKF_PREV_FORCED_N] = NAN;
+                        ost->forced_keyframes_expr_const_values[FKF_PREV_FORCED_T] = NAN;
+                    } else {
+                        parse_forced_key_frames(ost->forced_keyframes, ost, ost->st->codec);
+                    }
+                }
                 break;
             case AVMEDIA_TYPE_SUBTITLE:
                 codec->time_base = (AVRational){1, 1000};
diff --git a/ffmpeg.h b/ffmpeg.h
index 9e2e97b..c64a015 100644
--- a/ffmpeg.h
+++ b/ffmpeg.h
@@ -41,6 +41,7 @@
 
 #include "libavutil/avutil.h"
 #include "libavutil/dict.h"
+#include "libavutil/eval.h"
 #include "libavutil/fifo.h"
 #include "libavutil/pixfmt.h"
 #include "libavutil/rational.h"
@@ -289,6 +290,17 @@ typedef struct InputFile {
 #endif
 } InputFile;
 
+enum forced_keyframes_const {
+    FKF_N,
+    FKF_N_FORCED,
+    FKF_PREV_FORCED_N,
+    FKF_PREV_FORCED_T,
+    FKF_T,
+    FKF_NB
+};
+
+extern const char *const forced_keyframes_const_names[];
+
 typedef struct OutputStream {
     int file_index;          /* file index */
     int index;               /* stream index in the output file */
@@ -320,6 +332,8 @@ typedef struct OutputStream {
     int forced_kf_count;
     int forced_kf_index;
     char *forced_keyframes;
+    AVExpr *forced_keyframes_pexpr;
+    double forced_keyframes_expr_const_values[FKF_NB];
 
     /* audio only */
     int audio_channels_map[SWR_CH_MAX];  /* list of the channels id to pick from the source stream */