[FFmpeg-devel] [PATCH] Modified force_key_frames command line option to accept frame numbers

[Sylvester Zaluga]

 @item -force_key_frames[:@var{stream_specifier}] @var{time}[, at var{time}...] (@emph{output,per-stream})
+ at item -force_key_frames[:@var{stream_specifier}] n:@var{number}[, at var{number}...] (@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.
+Force key frames at the specified timestamps (more precisely at the first
+frames after each specified time) or at the specified frame numbers.
+
+If the argument is prefixed with @code{n:}, the comma-separated @var{number}
+strings will be interpreted as numbers / indices of frames that should
+be forced to be key frames.
 
 If the argument is prefixed with @code{expr:}, the string @var{expr}
 is interpreted like an expression and is evaluated for each frame. A
@@ -606,6 +611,11 @@ before the beginning of every chapter:
 -force_key_frames 0:05:00,chapters-0.1
 @end example
 
+For example, to insert at predefined frame numbers:
+ at example
+-force_key_frames n:0,25,50,75,100,125,150,175,200,225,250
+ at end example
+
 The expression in @var{expr} can contain the following constants:
 @table @option
 @item n
diff --git a/ffmpeg.c b/ffmpeg.c
index 1fd0ece..969222d 100644
--- a/ffmpeg.c
+++ b/ffmpeg.c
@@ -1039,6 +1039,11 @@ static void do_video_out(AVFormatContext *s,
             }
 
             ost->forced_keyframes_expr_const_values[FKF_N] += 1;
+        } else if (ost->forced_kf_n_index < ost->forced_kf_n_count &&
+            ost->frame_number == ost->forced_kf_n_pts[ost->forced_kf_n_index]) {
+            av_dlog(NULL, "force_key_frame (number mode): n:%d\n", ost->frame_number);
+            ost->forced_kf_n_index++;
+            forced_keyframe = 1;
         }
 
         if (forced_keyframe) {
@@ -2333,6 +2338,12 @@ static InputStream *get_input_stream(OutputStream *ost)
     return NULL;
 }
 
+static int compare_int(const void *a, const void *b)
+{
+    int va = *(int *)a, vb = *(int *)b;
+    return va < vb ? -1 : va > vb ? +1 : 0;
+}
+
 static int compare_int64(const void *a, const void *b)
 {
     int64_t va = *(int64_t *)a, vb = *(int64_t *)b;
@@ -2402,6 +2413,53 @@ static void parse_forced_key_frames(char *kf, OutputStream *ost,
     ost->forced_kf_pts   = pts;
 }
 
+static void parse_forced_key_frames_n(char *kf, OutputStream *ost,
+                                      AVCodecContext *avctx)
+{
+    char *p;
+    int n = 1, i, size, index = 0;
+    int kf_n, kf_n_prev, *pts;
+
+    for (p = kf; *p; p++)
+        if (*p == ',')
+            n++;
+    size = n;
+    pts = av_malloc_array(size, sizeof(*pts));
+    if (!pts) {
+        av_log(NULL, AV_LOG_FATAL, "Could not allocate forced key frames array.\n");
+        exit_program(1);
+    }
+
+    p = kf;
+    for (i = 0; i < n; i++) {
+        char *next = strchr(p, ',');
+
+        if (next)
+            *next++ = 0;
+
+        kf_n = parse_number_or_die("force_key_frames (numbers)", p, OPT_INT, 0, INT_MAX);
+        av_assert1(index < size);
+        pts[index++] = kf_n;
+
+        p = next;
+    }
+
+    av_assert0(index == size);
+    qsort(pts, size, sizeof(*pts), compare_int);
+    
+    kf_n_prev = -1;
+    for (i = 0; i < n; i++) {
+        if (pts[i] == kf_n_prev) {
+            av_log(NULL, AV_LOG_FATAL, "Duplicated forced key frame number.\n");
+            exit_program(1);
+        }
+        kf_n_prev = pts[i];
+    }
+        
+    ost->forced_kf_n_count = size;
+    ost->forced_kf_n_pts   = pts;
+}
+
 static void report_new_stream(int input_index, AVPacket *pkt)
 {
     InputFile *file = input_files[input_index];
@@ -2815,6 +2873,8 @@ static int transcode_init(void)
                         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 if (!strncmp(ost->forced_keyframes, "n:", 2)) {
+                        parse_forced_key_frames_n(ost->forced_keyframes+2, ost, ost->enc_ctx);
                     } else {
                         parse_forced_key_frames(ost->forced_keyframes, ost, ost->enc_ctx);
                     }
@@ -3775,6 +3835,7 @@ static int transcode(void)
                     ost->logfile = NULL;
                 }
                 av_freep(&ost->forced_kf_pts);
+                av_freep(&ost->forced_kf_n_pts);
                 av_freep(&ost->apad);
                 av_dict_free(&ost->encoder_opts);
                 av_dict_free(&ost->swr_opts);
diff --git a/ffmpeg.h b/ffmpeg.h
index c456603..539b8bd 100644
--- a/ffmpeg.h
+++ b/ffmpeg.h
@@ -400,6 +400,9 @@ typedef struct OutputStream {
     int64_t *forced_kf_pts;
     int forced_kf_count;
     int forced_kf_index;
+    int *forced_kf_n_pts;                /* list of forced key frame numbers / indices */
+    int forced_kf_n_count;               /* count of forced key frame numbers / indices */
+    int forced_kf_n_index;               /* current numbered forced key frame */
     char *forced_keyframes;
     AVExpr *forced_keyframes_pexpr;
     double forced_keyframes_expr_const_values[FKF_NB];
diff --git a/ffmpeg_opt.c b/ffmpeg_opt.c
index 77ef0c4..15c4813 100644
--- a/ffmpeg_opt.c
+++ b/ffmpeg_opt.c
@@ -2990,7 +2990,7 @@ const OptionDef options[] = {
         "set the value of an outfile streamid", "streamIndex:value" },
     { "force_key_frames", OPT_VIDEO | OPT_STRING | HAS_ARG | OPT_EXPERT |
                           OPT_SPEC | OPT_OUTPUT,                                 { .off = OFFSET(forced_key_frames) },
-        "force key frames at specified timestamps", "timestamps" },
+        "force key frames at specified timestamps or frame numbers", "timestamps/frame numbers" },
     { "ab",           OPT_VIDEO | HAS_ARG | OPT_PERFILE | OPT_OUTPUT,            { .func_arg = opt_bitrate },
         "audio bitrate (please use -b:a)", "bitrate" },
     { "b",            OPT_VIDEO | HAS_ARG | OPT_PERFILE | OPT_OUTPUT,            { .func_arg = opt_bitrate },