feat: expose ctx for muxer

This commit is contained in:
kieran 2024-11-14 10:34:36 +00:00
parent bde945fe88
commit aae40f0fba
No known key found for this signature in database
GPG Key ID: DE71CEB3925BE941

View File

@ -102,16 +102,15 @@ impl TryInto<*mut AVIOContext> for &mut MuxerOutput {
} }
pub struct MuxerBuilder { pub struct MuxerBuilder {
value: Muxer, ctx: *mut AVFormatContext,
output: MuxerOutput,
} }
impl MuxerBuilder { impl MuxerBuilder {
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
value: Muxer {
ctx: ptr::null_mut(), ctx: ptr::null_mut(),
output: MuxerOutput::Url(String::new()), output: MuxerOutput::Url(String::new()),
},
} }
} }
@ -121,12 +120,12 @@ impl MuxerBuilder {
format: Option<&str>, format: Option<&str>,
options: Option<HashMap<String, String>>, options: Option<HashMap<String, String>>,
) -> Result<()> { ) -> Result<()> {
if !self.value.ctx.is_null() { if !self.ctx.is_null() {
bail!("context already open"); bail!("context already open");
} }
let ret = avformat_alloc_output_context2( let ret = avformat_alloc_output_context2(
&mut self.value.ctx, &mut self.ctx,
ptr::null_mut(), ptr::null_mut(),
if let Some(format) = format { if let Some(format) = format {
cstr!(format) cstr!(format)
@ -142,13 +141,13 @@ impl MuxerBuilder {
bail_ffmpeg!(ret); bail_ffmpeg!(ret);
// Setup global header flag // Setup global header flag
if (*(*self.value.ctx).oformat).flags & AVFMT_GLOBALHEADER != 0 { if (*(*self.ctx).oformat).flags & AVFMT_GLOBALHEADER != 0 {
(*self.value.ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int; (*self.ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int;
} }
// Set options on ctx // Set options on ctx
if let Some(opts) = options { if let Some(opts) = options {
set_opts((*self.value.ctx).priv_data, opts)?; set_opts((*self.ctx).priv_data, opts)?;
} }
Ok(()) Ok(())
} }
@ -165,7 +164,7 @@ impl MuxerBuilder {
{ {
let path_str = dst.into(); let path_str = dst.into();
self.init_ctx(Some(path_str), format, options)?; self.init_ctx(Some(path_str), format, options)?;
self.value.output = MuxerOutput::Url(path_str.to_string()); self.output = MuxerOutput::Url(path_str.to_string());
Ok(self) Ok(self)
} }
@ -181,7 +180,7 @@ impl MuxerBuilder {
W: WriteSeek + 'static, W: WriteSeek + 'static,
{ {
self.init_ctx(None, format, options)?; self.init_ctx(None, format, options)?;
self.value.output = MuxerOutput::WriterSeeker(Some(slimbox_unsize!(writer))); self.output = MuxerOutput::WriterSeeker(Some(slimbox_unsize!(writer)));
Ok(self) Ok(self)
} }
@ -196,39 +195,50 @@ impl MuxerBuilder {
W: Write + 'static, W: Write + 'static,
{ {
self.init_ctx(None, format, options)?; self.init_ctx(None, format, options)?;
self.value.output = MuxerOutput::Writer(Some(slimbox_unsize!(writer))); self.output = MuxerOutput::Writer(Some(slimbox_unsize!(writer)));
Ok(self) Ok(self)
} }
/// Add a stream to the output using an existing encoder /// Add a stream to the output using an existing encoder
pub unsafe fn with_stream_encoder(mut self, encoder: &Encoder) -> Result<Self> { pub unsafe fn with_stream_encoder(self, encoder: &Encoder) -> Result<Self> {
self.value.add_stream_encoder(encoder)?; Self::add_stream_from_encoder(self.ctx, encoder)?;
Ok(self) Ok(self)
} }
/// Add a stream to the output using an existing input stream (copy) /// Add a stream to the output using an existing input stream (copy)
pub unsafe fn with_copy_stream(mut self, in_stream: *mut AVStream) -> Result<Self> { pub unsafe fn with_copy_stream(self, in_stream: *mut AVStream) -> Result<Self> {
self.value.add_copy_stream(in_stream)?; Self::add_copy_stream(self.ctx, in_stream)?;
Ok(self) Ok(self)
} }
/// Apply custom options to the [AVFormatContext]
pub unsafe fn with_custom_options<F>(self, f_mod: F) -> Self
where
F: FnOnce(*mut AVFormatContext),
{
f_mod(self.ctx);
self
}
/// Build the muxer /// Build the muxer
pub fn build(self) -> Result<Muxer> { pub fn build(self) -> Result<Muxer> {
if self.value.ctx.is_null() { if self.ctx.is_null() {
bail!("context is null"); bail!("context is null");
} }
Ok(self.value) Ok(Muxer {
} ctx: self.ctx,
} output: self.output,
})
impl Muxer {
pub fn builder() -> MuxerBuilder {
MuxerBuilder::new()
} }
/// Add a stream to the output using an existing encoder pub unsafe fn add_stream_from_encoder(
pub unsafe fn add_stream_encoder(&mut self, encoder: &Encoder) -> Result<*mut AVStream> { ctx: *mut AVFormatContext,
let stream = avformat_new_stream(self.ctx, encoder.codec()); encoder: &Encoder,
) -> Result<*mut AVStream> {
if ctx.is_null() {
bail!("cannot add stream to null ctx");
}
let stream = avformat_new_stream(ctx, encoder.codec());
if stream.is_null() { if stream.is_null() {
bail!("unable to allocate stream"); bail!("unable to allocate stream");
} }
@ -244,9 +254,14 @@ impl Muxer {
Ok(stream) Ok(stream)
} }
/// Add a stream to the output using an existing input stream (copy) pub(crate) unsafe fn add_copy_stream(
pub unsafe fn add_copy_stream(&mut self, in_stream: *mut AVStream) -> Result<*mut AVStream> { ctx: *mut AVFormatContext,
let stream = avformat_new_stream(self.ctx, ptr::null_mut()); in_stream: *mut AVStream,
) -> Result<*mut AVStream> {
if ctx.is_null() {
bail!("cannot add stream to null ctx");
}
let stream = avformat_new_stream(ctx, ptr::null_mut());
if stream.is_null() { if stream.is_null() {
bail!("unable to allocate stream"); bail!("unable to allocate stream");
} }
@ -257,6 +272,22 @@ impl Muxer {
Ok(stream) Ok(stream)
} }
}
impl Muxer {
pub fn builder() -> MuxerBuilder {
MuxerBuilder::new()
}
/// Add a stream to the output using an existing encoder
pub unsafe fn add_stream_encoder(&mut self, encoder: &Encoder) -> Result<*mut AVStream> {
MuxerBuilder::add_stream_from_encoder(self.ctx, encoder)
}
/// Add a stream to the output using an existing input stream (copy)
pub unsafe fn add_copy_stream(&mut self, in_stream: *mut AVStream) -> Result<*mut AVStream> {
MuxerBuilder::add_copy_stream(self.ctx, in_stream)
}
/// Open the output to start sending packets /// Open the output to start sending packets
pub unsafe fn open(&mut self) -> Result<()> { pub unsafe fn open(&mut self) -> Result<()> {
@ -278,6 +309,11 @@ impl Muxer {
Ok(()) Ok(())
} }
/// Get [AVFormatContext] pointer
pub fn context(&self) -> *mut AVFormatContext {
self.ctx
}
/// Write a packet to the output /// Write a packet to the output
pub unsafe fn write_packet(&mut self, pkt: *mut AVPacket) -> Result<()> { pub unsafe fn write_packet(&mut self, pkt: *mut AVPacket) -> Result<()> {
let stream = *(*self.ctx).streams.add((*pkt).stream_index as usize); let stream = *(*self.ctx).streams.add((*pkt).stream_index as usize);
@ -322,7 +358,6 @@ mod tests {
unsafe fn setup_encoder() -> Result<(*mut AVFrame, Encoder)> { unsafe fn setup_encoder() -> Result<(*mut AVFrame, Encoder)> {
std::fs::create_dir_all("test_output")?; std::fs::create_dir_all("test_output")?;
let path = PathBuf::from("test_output/test.mp4");
let frame = generate_test_frame(); let frame = generate_test_frame();
// convert frame to YUV // convert frame to YUV