set pts in ffmpeg screen encoder (#1020)

Author: Brendonovich

Cargo.lock

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use cap_media_info::{Pixel, VideoInfo};
 use ffmpeg::{
     Dictionary,
@@ -103,7 +105,7 @@ impl H264EncoderBuilder {
         encoder.set_width(input_config.width);
         encoder.set_height(input_config.height);
         encoder.set_format(format);
-        encoder.set_time_base(input_config.frame_rate.invert());
+        encoder.set_time_base(input_config.time_base);
         encoder.set_frame_rate(Some(input_config.frame_rate));
 
         // let target_bitrate = compression.bitrate();
@@ -121,7 +123,7 @@ impl H264EncoderBuilder {
 
         let mut output_stream = output.add_stream(codec)?;
         let stream_index = output_stream.index();
-        output_stream.set_time_base(input_config.frame_rate.invert());
+        // output_stream.set_time_base((1, H264Encoder::TIME_BASE));
         output_stream.set_rate(input_config.frame_rate);
         output_stream.set_parameters(&video_encoder);
 
@@ -147,6 +149,8 @@ pub struct H264Encoder {
 }
 
 impl H264Encoder {
+    const TIME_BASE: i32 = 15000;
+
     pub fn builder(name: &'static str, input_config: VideoInfo) -> H264EncoderBuilder {
         H264EncoderBuilder::new(name, input_config)
     }
@@ -203,6 +207,11 @@ impl H264Encoder {
         }
         self.process_frame(output);
     }
+
+    pub fn get_pts(&self, duration: Duration) -> i64 {
+        (duration.as_secs_f32() * self.config.time_base.denominator() as f32
+            / self.config.time_base.numerator() as f32) as i64
+    }
 }
 
 fn get_codec_and_options(

crates/enc-ffmpeg/src/video/h264.rs

@@ -56,7 +56,7 @@ pub async fn main() {
 
     tokio::time::sleep(Duration::from_millis(10)).await;
 
-    let (handle, _ready_rx) = instant_recording::Actor::builder(
+    let (handle, _ready_rx) = studio_recording::Actor::builder(
         dir.path().into(),
         ScreenCaptureTarget::Display {
             id: Display::primary().id(),

crates/recording/examples/recording-cli.rs

@@ -26,6 +26,7 @@ pub trait MakeCapturePipeline: ScreenCaptureFormat + std::fmt::Debug + 'static {
             flume::Receiver<(Self::VideoFormat, Timestamp)>,
         ),
         output_path: PathBuf,
+        start_time: Timestamps,
     ) -> Result<(PipelineBuilder, flume::Receiver<Timestamp>), MediaError>
     where
         Self: Sized;
@@ -54,6 +55,7 @@ impl MakeCapturePipeline for screen_capture::CMSampleBufferCapture {
             flume::Receiver<(Self::VideoFormat, Timestamp)>,
         ),
         output_path: PathBuf,
+        _start_time: Timestamps,
     ) -> Result<(PipelineBuilder, flume::Receiver<Timestamp>), MediaError> {
         let screen_config = source.0.info();
         tracing::info!("screen config: {:?}", screen_config);
@@ -234,6 +236,7 @@ impl MakeCapturePipeline for screen_capture::Direct3DCapture {
             flume::Receiver<(Self::VideoFormat, Timestamp)>,
         ),
         output_path: PathBuf,
+        start_time: Timestamps,
     ) -> Result<(PipelineBuilder, flume::Receiver<Timestamp>), MediaError>
     where
         Self: Sized,
@@ -279,10 +282,8 @@ impl MakeCapturePipeline for screen_capture::Direct3DCapture {
                     either::Left((encoder, muxer))
                 }
                 Err(e) => {
-                    use tracing::{error, info};
-
-                    error!("Failed to create native encoder: {e}");
-                    info!("Falling back to software H264 encoder");
+                    tracing::error!("Failed to create native encoder: {e}");
+                    tracing::info!("Falling back to software H264 encoder");
 
                     either::Right(
                         cap_enc_ffmpeg::H264Encoder::builder("screen", screen_config)
@@ -351,6 +352,7 @@ impl MakeCapturePipeline for screen_capture::Direct3DCapture {
                         .map_err(|e| format!("EncoderFinish: {e}"))?;
                 }
                 either::Right(mut encoder) => {
+                    let mut first_timestamp = None;
                     let mut timestamp_tx = Some(timestamp_tx);
                     let _ = ready.send(Ok(()));
 
@@ -361,10 +363,16 @@ impl MakeCapturePipeline for screen_capture::Direct3DCapture {
                             let _ = timestamp_tx.send(timestamp);
                         }
 
-                        let ff_frame = frame
+                        let first_timestamp = first_timestamp.get_or_insert(timestamp);
+
+                        let mut ff_frame = frame
                             .as_ffmpeg()
                             .map_err(|e| format!("FrameAsFfmpeg: {e}"))?;
 
+                        let elapsed = timestamp.duration_since(start_time)
+                            - first_timestamp.duration_since(start_time);
+                        ff_frame.set_pts(Some(encoder.get_pts(elapsed)));
+
                         encoder.queue_frame(ff_frame, &mut output);
                     }
                     encoder.finish(&mut output);

crates/recording/src/capture_pipeline.rs

@@ -752,6 +752,7 @@ async fn create_segment_pipeline(
                 pipeline_builder,
                 (screen_source, screen_rx),
                 screen_output_path.clone(),
+                start_time,
             )
             .unwrap();
         pipeline_builder = pipeline_builder_;