Enable faststart and fix seeking in outputted mp4s. (#897)

software-mansion · Dec 16, 2024 · 2505166 · 2505166
1 parent 0b74cc9
commit 2505166
Show file tree

Hide file tree

Showing 9 changed files with 46 additions and 8 deletions.
diff --git a/compositor_pipeline/src/pipeline/encoder/fdk_aac.rs b/compositor_pipeline/src/pipeline/encoder/fdk_aac.rs
@@ -14,7 +14,7 @@ use tracing::{debug, error, span, Level};
 use crate::{
     audio_mixer::{AudioChannels, AudioSamples, OutputSamples},
     error::EncoderInitError,
-    pipeline::{AudioCodec, EncodedChunk, EncodedChunkKind, EncoderOutputEvent},
+    pipeline::{types::IsKeyframe, AudioCodec, EncodedChunk, EncodedChunkKind, EncoderOutputEvent},
     queue::PipelineEvent,
 };
 
@@ -266,6 +266,7 @@ impl AacEncoderInner {
                 data: output.freeze(),
                 pts,
                 dts: None,
+                is_keyframe: IsKeyframe::NoKeyframes,
                 kind: EncodedChunkKind::Audio(AudioCodec::Aac),
             })),
         }

diff --git a/compositor_pipeline/src/pipeline/encoder/ffmpeg_h264.rs b/compositor_pipeline/src/pipeline/encoder/ffmpeg_h264.rs
@@ -13,7 +13,8 @@ use tracing::{debug, error, span, trace, warn, Level};
 use crate::{
     error::EncoderInitError,
     pipeline::types::{
-        ChunkFromFfmpegError, EncodedChunk, EncodedChunkKind, EncoderOutputEvent, VideoCodec,
+        ChunkFromFfmpegError, EncodedChunk, EncodedChunkKind, EncoderOutputEvent, IsKeyframe,
+        VideoCodec,
     },
     queue::PipelineEvent,
 };
@@ -389,6 +390,11 @@ fn encoded_chunk_from_av_packet(
             .map(rescale)
             .ok_or(ChunkFromFfmpegError::NoPts)?,
         dts: value.dts().map(rescale),
+        is_keyframe: if value.flags().contains(ffmpeg_next::packet::Flags::KEY) {
+            IsKeyframe::Yes
+        } else {
+            IsKeyframe::No
+        },
         kind,
     })
 }
diff --git a/compositor_pipeline/src/pipeline/encoder/opus.rs b/compositor_pipeline/src/pipeline/encoder/opus.rs
@@ -6,7 +6,7 @@ use crate::{
     audio_mixer::{AudioChannels, AudioSamples, OutputSamples},
     error::EncoderInitError,
     pipeline::{
-        types::{EncodedChunk, EncodedChunkKind, EncoderOutputEvent},
+        types::{EncodedChunk, EncodedChunkKind, EncoderOutputEvent, IsKeyframe},
         AudioCodec,
     },
     queue::PipelineEvent,
@@ -94,6 +94,7 @@ fn run_encoder_thread(
             data,
             pts: batch.start_pts,
             dts: None,
+            is_keyframe: IsKeyframe::NoKeyframes,
             kind: EncodedChunkKind::Audio(AudioCodec::Opus),
         };
 

diff --git a/compositor_pipeline/src/pipeline/input/mp4/reader.rs b/compositor_pipeline/src/pipeline/input/mp4/reader.rs
@@ -11,7 +11,7 @@ use mp4::Mp4Sample;
 use tracing::warn;
 
 use crate::pipeline::{
-    types::{EncodedChunk, EncodedChunkKind},
+    types::{EncodedChunk, EncodedChunkKind, IsKeyframe},
     AudioCodec, VideoCodec,
 };
 
@@ -235,6 +235,10 @@ impl<Reader: Read + Seek + Send + 'static> TrackChunks<'_, Reader> {
             data,
             pts,
             dts: Some(dts),
+            is_keyframe: match self.track.decoder_options {
+                DecoderOptions::H264 => IsKeyframe::Unknown,
+                DecoderOptions::Aac(_) => IsKeyframe::NoKeyframes,
+            },
             kind: match self.track.decoder_options {
                 DecoderOptions::H264 => EncodedChunkKind::Video(VideoCodec::H264),
                 DecoderOptions::Aac(_) => EncodedChunkKind::Audio(AudioCodec::Aac),

diff --git a/compositor_pipeline/src/pipeline/input/rtp/depayloader.rs b/compositor_pipeline/src/pipeline/input/rtp/depayloader.rs
@@ -10,7 +10,7 @@ use rtp::{
 use crate::pipeline::{
     decoder::{self, AacDecoderOptions},
     rtp::{AUDIO_PAYLOAD_TYPE, VIDEO_PAYLOAD_TYPE},
-    types::{AudioCodec, EncodedChunk, EncodedChunkKind, VideoCodec},
+    types::{AudioCodec, EncodedChunk, EncodedChunkKind, IsKeyframe, VideoCodec},
     VideoDecoder,
 };
 
@@ -126,6 +126,7 @@ impl VideoDepayloader {
                     data: mem::take(buffer).concat().into(),
                     pts: Duration::from_secs_f64(timestamp as f64 / 90000.0),
                     dts: None,
+                    is_keyframe: IsKeyframe::Unknown,
                     kind,
                 };
 
@@ -194,6 +195,7 @@ impl AudioDepayloader {
                     data: opus_packet,
                     pts: Duration::from_secs_f64(timestamp as f64 / 48000.0),
                     dts: None,
+                    is_keyframe: IsKeyframe::NoKeyframes,
                     kind,
                 }])
             }

diff --git a/compositor_pipeline/src/pipeline/input/rtp/depayloader/aac.rs b/compositor_pipeline/src/pipeline/input/rtp/depayloader/aac.rs
@@ -4,7 +4,7 @@ use bytes::{Buf, BytesMut};
 
 use crate::pipeline::{
     decoder::AacDepayloaderMode,
-    types::{EncodedChunk, EncodedChunkKind},
+    types::{EncodedChunk, EncodedChunkKind, IsKeyframe},
     AudioCodec,
 };
 
@@ -269,6 +269,7 @@ impl AacDepayloader {
                 pts,
                 data: payload,
                 dts: None,
+                is_keyframe: IsKeyframe::NoKeyframes,
                 kind: EncodedChunkKind::Audio(AudioCodec::Aac),
             });
         }

diff --git a/compositor_pipeline/src/pipeline/output/mp4.rs b/compositor_pipeline/src/pipeline/output/mp4.rs
@@ -10,7 +10,10 @@ use crate::{
     audio_mixer::AudioChannels,
     error::OutputInitError,
     event::Event,
-    pipeline::{EncodedChunk, EncodedChunkKind, EncoderOutputEvent, PipelineCtx, VideoCodec},
+    pipeline::{
+        types::IsKeyframe, EncodedChunk, EncodedChunkKind, EncoderOutputEvent, PipelineCtx,
+        VideoCodec,
+    },
 };
 
 #[derive(Debug, Clone)]
@@ -183,8 +186,10 @@ fn init_ffmpeg_output(
         })
         .transpose()?;
 
+    let ffmpeg_options = ffmpeg::Dictionary::from_iter(&[("movflags", "faststart")]);
+
     output_ctx
-        .write_header()
+        .write_header_with(ffmpeg_options)
         .map_err(OutputInitError::FfmpegMp4Error)?;
 
     Ok((output_ctx, video_stream, audio_stream))
@@ -287,6 +292,12 @@ fn create_packet(
     packet.set_time_base(ffmpeg::Rational::new(1, stream_state.time_base as i32));
     packet.set_stream(stream_state.id);
 
+    match chunk.is_keyframe {
+        IsKeyframe::Yes => packet.set_flags(ffmpeg::packet::Flags::KEY),
+        IsKeyframe::Unknown => warn!("The MP4 output received an encoded chunk with is_keyframe set to Unknown. This output needs this information to produce correct mp4s."),
+        IsKeyframe::NoKeyframes | IsKeyframe::No => {},
+    }
+
     Some(packet)
 }
 

diff --git a/compositor_pipeline/src/pipeline/types.rs b/compositor_pipeline/src/pipeline/types.rs
@@ -17,9 +17,21 @@ pub struct EncodedChunk {
     pub data: Bytes,
     pub pts: Duration,
     pub dts: Option<Duration>,
+    pub is_keyframe: IsKeyframe,
     pub kind: EncodedChunkKind,
 }
 
+pub enum IsKeyframe {
+    /// this is a keyframe
+    Yes,
+    /// this is not a keyframe
+    No,
+    /// it's unknown whether this frame is a keyframe or not
+    Unknown,
+    /// the codec this chunk is encoded in does not have keyframes at all
+    NoKeyframes,
+}
+
 #[derive(Debug)]
 pub enum EncoderOutputEvent {
     Data(EncodedChunk),

diff --git a/integration_tests/examples/mp4.rs → integration_tests/examples/mp4_output.rs b/integration_tests/examples/mp4.rs → integration_tests/examples/mp4_output.rs