Skip to main content

lumen_ffmpeg/encode/
muxed.rs

1use crate::audio::AudioFrame;
2use crate::gpu::GpuVideoInput;
3use crate::video::CpuVideoFrame;
4use crate::{FfmpegError, Result};
5
6use super::audio::{AudioEncoder, AudioEncoderConfig};
7use super::output::OutputContext;
8use super::telemetry::GpuEncodeTelemetry;
9use super::video::{VideoEncoder, VideoEncoderConfig};
10
11pub struct MuxedEncoder {
12    output: OutputContext,
13    video: VideoEncoder,
14    audio: Option<AudioEncoder>,
15    wrote_header: bool,
16}
17
18impl MuxedEncoder {
19    pub fn create(path: impl Into<String>, video: VideoEncoderConfig) -> Result<Self> {
20        Self::create_with_audio(path, video, None)
21    }
22
23    pub fn create_with_audio(
24        path: impl Into<String>,
25        video: VideoEncoderConfig,
26        audio: Option<AudioEncoderConfig>,
27    ) -> Result<Self> {
28        let mut output = OutputContext::create(path)?;
29        let mut video = VideoEncoder::create(&mut output, video)?;
30        let mut audio = audio
31            .map(|config| AudioEncoder::create(&mut output, config))
32            .transpose()?;
33        output.write_header()?;
34        video.refresh_stream_time_base(&output)?;
35        if let Some(audio) = audio.as_mut() {
36            audio.refresh_stream_time_base(&output)?;
37        }
38        Ok(Self {
39            output,
40            video,
41            audio,
42            wrote_header: true,
43        })
44    }
45
46    pub fn write_video_frame(&mut self, frame: &CpuVideoFrame) -> Result<()> {
47        self.video.send_cpu_frame(&mut self.output, frame)
48    }
49
50    pub fn write_gpu_frame(&mut self, frame: &GpuVideoInput<'_>) -> Result<()> {
51        self.video.send_gpu_frame(&mut self.output, frame)
52    }
53
54    pub fn write_audio_frame(&mut self, frame: &AudioFrame) -> Result<()> {
55        let Some(audio) = self.audio.as_mut() else {
56            return Err(FfmpegError::new(
57                "MuxedEncoder::write_audio_frame",
58                "encoder was created without an audio stream",
59            ));
60        };
61        audio.send_audio_frame(&mut self.output, frame)
62    }
63
64    pub fn gpu_telemetry(&self) -> &GpuEncodeTelemetry {
65        self.video.gpu_telemetry()
66    }
67
68    pub fn finish(mut self) -> Result<()> {
69        self.video.flush(&mut self.output)?;
70        if let Some(audio) = self.audio.as_mut() {
71            audio.flush(&mut self.output)?;
72        }
73        self.wrote_header = false;
74        self.output.write_trailer()
75    }
76}