From 62621bd38058f42303fe77da721696a3fe460af7 Mon Sep 17 00:00:00 2001
From: Romain Beauxis <toots@rastageeks.org>
Date: Sun, 26 Mar 2023 16:27:07 -0500
Subject: [PATCH] Fix PTS discontinuities in raw ffmpeg content. (#2978)

---
 src/encoder/ffmpeg_internal_encoder.ml |  9 +++++---
 src/tools/ffmpeg_avfilter_utils.ml     | 30 ++++++++++++++++++++------
 src/tools/ffmpeg_avfilter_utils.mli    |  1 +
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/encoder/ffmpeg_internal_encoder.ml b/src/encoder/ffmpeg_internal_encoder.ml
index ed9d4d67f0..f34ff53834 100644
--- a/src/encoder/ffmpeg_internal_encoder.ml
+++ b/src/encoder/ffmpeg_internal_encoder.ml
@@ -339,10 +339,13 @@ let mk_video ~ffmpeg ~options output =
 
   let converter = ref None in
 
+  let start_pts = ref 0L in
+
   let mk_converter ~pixel_format ~time_base ~stream_idx () =
     let c =
-      Ffmpeg_avfilter_utils.Fps.init ~width:target_width ~height:target_height
-        ~pixel_format ~time_base ~pixel_aspect ~target_fps ()
+      Ffmpeg_avfilter_utils.Fps.init ~start_pts:!start_pts ~width:target_width
+        ~height:target_height ~pixel_format ~time_base ~pixel_aspect ~target_fps
+        ()
     in
     converter := Some (pixel_format, time_base, stream_idx, c);
     c
@@ -353,7 +356,6 @@ let mk_video ~ffmpeg ~options output =
       | None -> mk_converter ~stream_idx ~pixel_format ~time_base ()
       | Some (p, t, i, _) when (p, t, i) <> (pixel_format, time_base, stream_idx)
         ->
-          log#important "Frame format change detected!";
           mk_converter ~stream_idx ~pixel_format ~time_base ()
       | Some (_, _, _, c) -> c
   in
@@ -376,6 +378,7 @@ let mk_video ~ffmpeg ~options output =
             (Ffmpeg_utils.best_pts frame)
         in
         Avutil.Frame.set_pts frame frame_pts;
+        start_pts := Int64.succ !start_pts;
         Av.write_frame stream frame)
   in
 
diff --git a/src/tools/ffmpeg_avfilter_utils.ml b/src/tools/ffmpeg_avfilter_utils.ml
index 9d0b0b14f2..310456e3b4 100644
--- a/src/tools/ffmpeg_avfilter_utils.ml
+++ b/src/tools/ffmpeg_avfilter_utils.ml
@@ -33,8 +33,8 @@ module Fps = struct
     | `Filter { time_base } -> time_base
     | `Pass_through time_base -> time_base
 
-  let init ~width ~height ~pixel_format ~time_base ?pixel_aspect ?source_fps
-      ~target_fps () =
+  let init ?(start_pts = 0L) ~width ~height ~pixel_format ~time_base
+      ?pixel_aspect ?source_fps ~target_fps () =
     let config = Avfilter.init () in
     let _buffer =
       let args =
@@ -70,6 +70,21 @@ module Fps = struct
       in
       Avfilter.attach ~name:"fps" ~args fps config
     in
+    let setpts =
+      match
+        List.find_opt
+          (fun { Avfilter.name } -> name = "setpts")
+          Avfilter.filters
+      with
+        | Some setpts -> setpts
+        | None -> failwith "Could not find setpts ffmpeg filter!"
+    in
+    let setpts =
+      let args =
+        [`Pair ("expr", `String (Printf.sprintf "%Ld+PTS-STARTPTS" start_pts))]
+      in
+      Avfilter.attach ~name:"setpts" ~args setpts config
+    in
     let _buffersink =
       Avfilter.attach ~name:"buffersink" Avfilter.buffersink config
     in
@@ -78,6 +93,9 @@ module Fps = struct
       (List.hd Avfilter.(fps.io.inputs.video));
     Avfilter.link
       (List.hd Avfilter.(fps.io.outputs.video))
+      (List.hd Avfilter.(setpts.io.inputs.video));
+    Avfilter.link
+      (List.hd Avfilter.(setpts.io.outputs.video))
       (List.hd Avfilter.(_buffersink.io.inputs.video));
     let graph = Avfilter.launch config in
     let _, input = List.hd Avfilter.(graph.inputs.video) in
@@ -86,14 +104,14 @@ module Fps = struct
     { input; output; time_base }
 
   (* Source fps is not always known so it is optional here. *)
-  let init ~width ~height ~pixel_format ~time_base ?pixel_aspect ?source_fps
-      ~target_fps () =
+  let init ?start_pts ~width ~height ~pixel_format ~time_base ?pixel_aspect
+      ?source_fps ~target_fps () =
     match source_fps with
       | Some f when f = target_fps -> `Pass_through time_base
       | _ ->
           `Filter
-            (init ~width ~height ~pixel_format ~time_base ?pixel_aspect
-               ?source_fps ~target_fps ())
+            (init ?start_pts ~width ~height ~pixel_format ~time_base
+               ?pixel_aspect ?source_fps ~target_fps ())
 
   let convert converter frame cb =
     match converter with
diff --git a/src/tools/ffmpeg_avfilter_utils.mli b/src/tools/ffmpeg_avfilter_utils.mli
index 82c8daa9bd..f36ca10e3d 100644
--- a/src/tools/ffmpeg_avfilter_utils.mli
+++ b/src/tools/ffmpeg_avfilter_utils.mli
@@ -26,6 +26,7 @@ module Fps : sig
   val time_base : t -> Avutil.rational
 
   val init :
+    ?start_pts:int64 ->
     width:int ->
     height:int ->
     pixel_format:Avutil.Pixel_format.t ->