From 8032966982ce6d79302a7ffed1e2d064a7ba5eb0 Mon Sep 17 00:00:00 2001 From: kieran Date: Sat, 9 Nov 2024 15:47:03 +0000 Subject: [PATCH] feat: muxer --- .gitignore | 5 +- examples/main.rs | 2 +- src/decode.rs | 28 ++++--- src/demux.rs | 4 +- src/encode.rs | 56 +++++++++----- src/filter.rs | 16 ++-- src/lib.rs | 21 ++++-- src/muxer.rs | 185 +++++++++++++++++++++++++++++++++++++++++++++++ src/resample.rs | 20 +++-- src/scale.rs | 13 ++-- src/transcode.rs | 9 +++ 11 files changed, 299 insertions(+), 60 deletions(-) create mode 100644 src/muxer.rs create mode 100644 src/transcode.rs diff --git a/.gitignore b/.gitignore index adf1101..86b2d2e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,7 @@ /target /.idea -test.png \ No newline at end of file +# Test output +*.mp4 +*.png +*.mkv \ No newline at end of file diff --git a/examples/main.rs b/examples/main.rs index 9a098d6..537847a 100644 --- a/examples/main.rs +++ b/examples/main.rs @@ -9,7 +9,7 @@ use std::path::PathBuf; fn main() { env_logger::init(); let name = args().next().unwrap_or("main".to_string()); - let path = if let Some(path) = args().skip(1).next() { + let path = if let Some(path) = args().nth(1) { PathBuf::from(path) } else { error!("Usage: {} ", name); diff --git a/src/decode.rs b/src/decode.rs index 744c675..ec21b46 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -1,4 +1,4 @@ -use crate::{options_to_dict, return_ffmpeg_error, rstr, StreamInfoChannel}; +use crate::{bail_ffmpeg, options_to_dict, rstr, StreamInfoChannel}; use std::collections::hash_map::Entry; use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; @@ -35,9 +35,9 @@ impl DecoderCodecContext { /// Get the codec name pub fn codec_name(&self) -> String { - let codec_name = unsafe { rstr!((*(*self).codec).name) }; + let codec_name = unsafe { rstr!((*self.codec).name) }; if self.hw_config.is_null() { - format!("{}", codec_name) + codec_name.to_string() } else { let hw = unsafe { rstr!(av_hwdevice_get_type_name((*self.hw_config).device_type)) }; format!("{}_{}", codec_name, hw) @@ -48,9 +48,9 @@ impl DecoderCodecContext { impl Drop for DecoderCodecContext { fn drop(&mut self) { unsafe { - avcodec_free_context(&mut self.context); - self.codec = ptr::null_mut(); - self.context = ptr::null_mut(); + if !self.context.is_null() { + avcodec_free_context(&mut self.context); + } } } } @@ -66,8 +66,6 @@ impl Display for DecoderCodecContext { } } -unsafe impl Send for DecoderCodecContext {} - pub struct Decoder { /// Decoder instances by stream index codecs: HashMap, @@ -75,6 +73,12 @@ pub struct Decoder { hw_decoder_types: Option>, } +impl Default for Decoder { + fn default() -> Self { + Self::new() + } +} + impl Decoder { pub fn new() -> Self { Self { @@ -177,7 +181,7 @@ impl Decoder { } let mut ret = avcodec_parameters_to_context(context, (*stream).codecpar); - return_ffmpeg_error!(ret, "Failed to copy codec parameters to context"); + bail_ffmpeg!(ret, "Failed to copy codec parameters to context"); let codec_name = rstr!(avcodec_get_name((*codec).id)); // try use HW decoder @@ -205,7 +209,7 @@ impl Decoder { ptr::null_mut(), 0, ); - return_ffmpeg_error!(ret, "Failed to create HW ctx"); + bail_ffmpeg!(ret, "Failed to create HW ctx"); (*context).hw_device_ctx = av_buffer_ref(hw_buf_ref); break; } @@ -218,7 +222,7 @@ impl Decoder { }; ret = avcodec_open2(context, codec, &mut dict); - return_ffmpeg_error!(ret, "Failed to open codec"); + bail_ffmpeg!(ret, "Failed to open codec"); let ctx = DecoderCodecContext { context, @@ -252,7 +256,7 @@ impl Decoder { stream: *mut AVStream, ) -> Result, Error> { let mut ret = avcodec_send_packet(ctx, pkt); - return_ffmpeg_error!(ret, "Failed to decode packet"); + bail_ffmpeg!(ret, "Failed to decode packet"); let mut pkgs = Vec::new(); while ret >= 0 { diff --git a/src/demux.rs b/src/demux.rs index 1199776..85cc9c2 100644 --- a/src/demux.rs +++ b/src/demux.rs @@ -1,4 +1,4 @@ -use crate::{cstr, return_ffmpeg_error}; +use crate::{bail_ffmpeg, cstr}; use crate::{get_ffmpeg_error_msg, DemuxerInfo, StreamChannelType, StreamInfoChannel}; use anyhow::Error; use ffmpeg_sys_the_third::*; @@ -102,7 +102,7 @@ impl Demuxer { pub unsafe fn probe_input(&mut self) -> Result { let ret = self.open_input(); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); if avformat_find_stream_info(self.ctx, ptr::null_mut()) < 0 { return Err(Error::msg("Could not find stream info")); diff --git a/src/encode.rs b/src/encode.rs index c574061..967b123 100644 --- a/src/encode.rs +++ b/src/encode.rs @@ -1,13 +1,14 @@ use std::collections::HashMap; use std::{ptr, slice}; -use crate::{get_ffmpeg_error_msg, options_to_dict, return_ffmpeg_error}; +use crate::{bail_ffmpeg, get_ffmpeg_error_msg, options_to_dict}; use anyhow::{bail, Error}; use ffmpeg_sys_the_third::{ - av_channel_layout_default, av_packet_alloc, av_packet_free, avcodec_alloc_context3, - avcodec_find_encoder, avcodec_get_supported_config, avcodec_open2, avcodec_receive_packet, - avcodec_send_frame, AVChannelLayout, AVCodec, AVCodecConfig, AVCodecContext, AVCodecID, - AVFrame, AVPacket, AVPixelFormat, AVRational, AVSampleFormat, AVERROR, AVERROR_EOF, + av_channel_layout_default, av_d2q, av_inv_q, av_packet_alloc, av_packet_free, + avcodec_alloc_context3, avcodec_find_encoder, avcodec_free_context, + avcodec_get_supported_config, avcodec_open2, avcodec_receive_packet, avcodec_send_frame, + AVChannelLayout, AVCodec, AVCodecConfig, AVCodecContext, AVCodecID, AVFrame, AVPacket, + AVPixelFormat, AVRational, AVSampleFormat, AVERROR, AVERROR_EOF, }; use libc::EAGAIN; @@ -16,7 +17,15 @@ pub struct Encoder { codec: *const AVCodec, } -unsafe impl Send for Encoder {} +impl Drop for Encoder { + fn drop(&mut self) { + unsafe { + if !self.ctx.is_null() { + avcodec_free_context(&mut self.ctx); + } + } + } +} impl Encoder { /// Create a new encoder with the specified codec @@ -36,6 +45,16 @@ impl Encoder { } } + /// Get the codec + pub fn codec(&self) -> *const AVCodec { + self.codec + } + + /// Get the codec context + pub fn codec_context(&self) -> *const AVCodecContext { + self.ctx + } + /// List supported configs (see [avcodec_get_supported_config]) pub unsafe fn list_configs<'a, T>(&mut self, cfg: AVCodecConfig) -> Result<&'a [T], Error> { let mut dst = ptr::null_mut(); @@ -48,16 +67,10 @@ impl Encoder { ptr::addr_of_mut!(dst) as _, &mut num_dst, ); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); Ok(slice::from_raw_parts(dst, num_dst as usize)) } - /// Set the encoder timebase - pub unsafe fn with_timebase(self, timebase: AVRational) -> Self { - (*self.ctx).time_base = timebase; - self - } - /// Set the encoder bitrate pub unsafe fn with_bitrate(self, bitrate: i64) -> Self { (*self.ctx).bit_rate = bitrate; @@ -94,9 +107,11 @@ impl Encoder { self } - /// Set the encoder profile (see AV_PROFILE_*) - pub unsafe fn with_framerate(self, rate: AVRational) -> Self { - (*self.ctx).framerate = rate; + /// Set the encoder framerate + pub unsafe fn with_framerate(self, fps: f32) -> Self { + let q = av_d2q(fps as f64, 90_000); + (*self.ctx).framerate = q; + (*self.ctx).time_base = av_inv_q(q); self } @@ -137,13 +152,15 @@ impl Encoder { /// Open the encoder so that you can start encoding frames (see [avcodec_open2]) pub unsafe fn open(self, options: Option>) -> Result { + assert!(!self.ctx.is_null()); + let mut options = if let Some(options) = options { options_to_dict(options)? } else { ptr::null_mut() }; let ret = avcodec_open2(self.ctx, self.codec, &mut options); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); Ok(self) } @@ -169,6 +186,7 @@ impl Encoder { } bail!(get_ffmpeg_error_msg(ret)); } + (*pkt).time_base = (*self.ctx).time_base; pkgs.push(pkt); } @@ -187,8 +205,8 @@ mod tests { unsafe { let frame = generate_test_frame(); let mut encoder = Encoder::new(AVCodecID::AV_CODEC_ID_PNG)? - .with_width(512) - .with_height(512); + .with_width((*frame).width) + .with_height((*frame).height); let pix_fmts: &[AVPixelFormat] = encoder.list_configs(AVCodecConfig::AV_CODEC_CONFIG_PIX_FORMAT)?; diff --git a/src/filter.rs b/src/filter.rs index 021e342..0e861cb 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -1,4 +1,4 @@ -use crate::{cstr, return_ffmpeg_error, rstr, set_opts}; +use crate::{bail_ffmpeg, cstr, rstr, set_opts}; use anyhow::Error; use ffmpeg_sys_the_third::{ av_strdup, avfilter_get_by_name, avfilter_graph_alloc, avfilter_graph_alloc_filter, @@ -13,6 +13,12 @@ pub struct Filter { graph: *mut AVFilterGraph, } +impl Default for Filter { + fn default() -> Self { + Self::new() + } +} + impl Filter { pub fn new() -> Self { Self { @@ -39,7 +45,7 @@ impl Filter { ptr::null_mut(), ctx, ); - return_ffmpeg_error!(ret, "Failed to parse graph"); + bail_ffmpeg!(ret, "Failed to parse graph"); let ret = avfilter_graph_create_filter( &mut dst_ctx, @@ -49,7 +55,7 @@ impl Filter { ptr::null_mut(), ctx, ); - return_ffmpeg_error!(ret, "Failed to parse graph"); + bail_ffmpeg!(ret, "Failed to parse graph"); (*outputs).name = av_strdup((*dst).name); (*outputs).filter_ctx = dst_ctx; @@ -62,7 +68,7 @@ impl Filter { (*inputs).next = ptr::null_mut(); let ret = avfilter_graph_parse(ctx, cstr!(graph), inputs, outputs, ptr::null_mut()); - return_ffmpeg_error!(ret, "Failed to parse graph"); + bail_ffmpeg!(ret, "Failed to parse graph"); let mut ret = Self { graph: ctx }; ret.build()?; Ok(ret) @@ -97,7 +103,7 @@ impl Filter { debug!("{}", d); let ret = avfilter_graph_config(self.graph, ptr::null_mut()); - return_ffmpeg_error!(ret, "Failed to build filter"); + bail_ffmpeg!(ret, "Failed to build filter"); Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index 3e03e84..98cc6e7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,20 +11,22 @@ mod decode; mod demux; mod encode; mod filter; +mod muxer; mod resample; mod scale; mod stream_info; +mod transcode; #[macro_export] -macro_rules! return_ffmpeg_error { +macro_rules! bail_ffmpeg { ($x:expr) => { if $x < 0 { - anyhow::bail!(crate::get_ffmpeg_error_msg($x)) + anyhow::bail!($crate::get_ffmpeg_error_msg($x)) } }; ($x:expr,$msg:expr) => { if $x < 0 { - anyhow::bail!(format!("{}: {}", $msg, crate::get_ffmpeg_error_msg($x))) + anyhow::bail!(format!("{}: {}", $msg, $crate::get_ffmpeg_error_msg($x))) } }; } @@ -60,7 +62,7 @@ unsafe fn options_to_dict(options: HashMap) -> Result<*mut AVDic let mut dict = ptr::null_mut(); for (key, value) in options { let ret = av_dict_set(&mut dict, cstr!(key), cstr!(value), 0); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); } Ok(dict) } @@ -109,7 +111,7 @@ fn set_opts(ctx: *mut libc::c_void, options: HashMap) -> Result< unsafe { for (key, value) in options { let ret = av_opt_set(ctx, cstr!(key), cstr!(value), AV_OPT_SEARCH_CHILDREN); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); } } Ok(()) @@ -122,7 +124,7 @@ pub unsafe fn get_frame_from_hw(mut frame: *mut AVFrame) -> Result<*mut AVFrame, } else { let new_frame = av_frame_alloc(); let ret = av_hwframe_transfer_data(new_frame, frame, 0); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); av_frame_copy_props(new_frame, frame); av_frame_free(&mut frame); Ok(new_frame) @@ -135,8 +137,8 @@ pub unsafe fn generate_test_frame() -> *mut AVFrame { use std::mem::transmute; let frame = av_frame_alloc(); - (*frame).width = 512; - (*frame).height = 512; + (*frame).width = 1024; + (*frame).height = 1024; (*frame).format = transmute(AVPixelFormat::AV_PIX_FMT_RGB24); av_frame_get_buffer(frame, 0); @@ -161,8 +163,11 @@ pub unsafe fn generate_test_frame() -> *mut AVFrame { pub use decode::*; pub use demux::*; +pub use encode::*; pub use ffmpeg_sys_the_third; pub use filter::*; +pub use muxer::*; pub use resample::*; pub use scale::*; pub use stream_info::*; +pub use transcode::*; diff --git a/src/muxer.rs b/src/muxer.rs new file mode 100644 index 0000000..3976607 --- /dev/null +++ b/src/muxer.rs @@ -0,0 +1,185 @@ +use crate::{bail_ffmpeg, cstr, set_opts, Encoder}; +use anyhow::{bail, Result}; +use ffmpeg_sys_the_third::{ + av_dump_format, av_interleaved_write_frame, av_packet_rescale_ts, av_write_trailer, + avcodec_parameters_from_context, avformat_alloc_output_context2, avformat_free_context, + avformat_new_stream, avformat_write_header, avio_flush, avio_open, AVFormatContext, AVPacket, + AVFMT_GLOBALHEADER, AVFMT_NOFILE, AVIO_FLAG_WRITE, AV_CODEC_FLAG_GLOBAL_HEADER, +}; +use std::collections::HashMap; +use std::path::PathBuf; +use std::ptr; + +pub struct Muxer { + ctx: *mut AVFormatContext, +} + +impl Drop for Muxer { + fn drop(&mut self) { + unsafe { + if !self.ctx.is_null() { + avio_flush((*self.ctx).pb); + avformat_free_context(self.ctx); + } + } + } +} + +impl Default for Muxer { + fn default() -> Self { + Self::new() + } +} + +impl Muxer { + pub fn new() -> Self { + Self { + ctx: ptr::null_mut(), + } + } + + /// Open the muxer with a destination path + pub unsafe fn with_output( + mut self, + dst: &PathBuf, + format: Option<&str>, + options: Option>, + ) -> Result { + if !self.ctx.is_null() { + bail!("context already open"); + } + + let ret = avformat_alloc_output_context2( + &mut self.ctx, + ptr::null_mut(), + if let Some(ref format) = format { + cstr!(format) + } else { + ptr::null() + }, + cstr!(dst.to_str().unwrap()), + ); + bail_ffmpeg!(ret); + + // Setup global header flag + if (*(*self.ctx).oformat).flags & AVFMT_GLOBALHEADER != 0 { + (*self.ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int; + } + + // Set options on ctx + if let Some(opts) = options { + set_opts((*self.ctx).priv_data, opts)?; + } + + Ok(self) + } + + /// Add a stream to the output using an existing encoder + pub unsafe fn with_stream_encoder(self, encoder: &Encoder) -> Result { + let stream = avformat_new_stream(self.ctx, encoder.codec()); + if stream.is_null() { + bail!("unable to allocate stream"); + } + let ret = avcodec_parameters_from_context((*stream).codecpar, encoder.codec_context()); + bail_ffmpeg!(ret); + + // setup other stream params + let encoder_ctx = encoder.codec_context(); + (*stream).time_base = (*encoder_ctx).time_base; + (*stream).avg_frame_rate = (*encoder_ctx).framerate; + (*stream).r_frame_rate = (*encoder_ctx).framerate; + + Ok(self) + } + + /// Open the output to start sending packets + pub unsafe fn open(self) -> Result { + if (*(*self.ctx).oformat).flags & AVFMT_NOFILE == 0 { + let ret = avio_open(&mut (*self.ctx).pb, (*self.ctx).url, AVIO_FLAG_WRITE); + bail_ffmpeg!(ret); + } + + let ret = avformat_write_header(self.ctx, ptr::null_mut()); + bail_ffmpeg!(ret); + + av_dump_format(self.ctx, 0, (*self.ctx).url, 1); + + Ok(self) + } + + /// Write a packet to the output + pub unsafe fn write_packet(&mut self, pkt: *mut AVPacket) -> Result<()> { + assert!((*pkt).stream_index >= 0 && (*pkt).stream_index < (*self.ctx).nb_streams as i32); + assert!((*pkt).time_base.num != 0 && (*pkt).time_base.den != 0); + + let stream = *(*self.ctx).streams.add((*pkt).stream_index as usize); + av_packet_rescale_ts(pkt, (*pkt).time_base, (*stream).time_base); + + let ret = av_interleaved_write_frame(self.ctx, pkt); + bail_ffmpeg!(ret); + Ok(()) + } + + /// Close the output and write the trailer + pub unsafe fn close(self) -> Result<()> { + let ret = av_write_trailer(self.ctx); + bail_ffmpeg!(ret); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{generate_test_frame, Scaler}; + use ffmpeg_sys_the_third::AVCodecID::AV_CODEC_ID_H264; + use ffmpeg_sys_the_third::AVPixelFormat::AV_PIX_FMT_YUV420P; + use ffmpeg_sys_the_third::AV_PROFILE_H264_MAIN; + + #[test] + fn encode_mkv() -> Result<()> { + unsafe { + let path = PathBuf::from("test.mp4"); + let frame = generate_test_frame(); + + // convert frame to YUV + let mut scaler = Scaler::new(); + let frame = scaler.process_frame( + frame, + (*frame).width as u16, + (*frame).height as u16, + AV_PIX_FMT_YUV420P, + )?; + + let mut encoder = Encoder::new(AV_CODEC_ID_H264)? + .with_width((*frame).width) + .with_height((*frame).height) + .with_pix_fmt(AV_PIX_FMT_YUV420P) + .with_bitrate(1_000_000) + .with_framerate(30.0) + .with_profile(AV_PROFILE_H264_MAIN) + .with_level(50) + .open(None)?; + + let mut muxer = Muxer::new() + .with_output(&path, None, None)? + .with_stream_encoder(&encoder)? + .open()?; + + let mut pts = 0; + for z in 0..100 { + (*frame).pts = pts; + for pkt in encoder.encode_frame(frame)? { + muxer.write_packet(pkt)?; + } + pts += 1; + } + // flush + for f_pk in encoder.encode_frame(ptr::null_mut())? { + muxer.write_packet(f_pk)?; + } + muxer.close()?; + } + Ok(()) + } +} diff --git a/src/resample.rs b/src/resample.rs index 21be9f8..d31853e 100644 --- a/src/resample.rs +++ b/src/resample.rs @@ -1,10 +1,10 @@ +use crate::bail_ffmpeg; use crate::get_ffmpeg_error_msg; -use crate::return_ffmpeg_error; use anyhow::Error; use ffmpeg_sys_the_third::{ av_channel_layout_default, av_frame_alloc, av_frame_copy_props, av_frame_free, - swr_alloc_set_opts2, swr_convert_frame, swr_init, AVChannelLayout, AVFrame, AVSampleFormat, - SwrContext, + swr_alloc_set_opts2, swr_convert_frame, swr_free, swr_init, AVChannelLayout, AVFrame, + AVSampleFormat, SwrContext, }; use libc::malloc; use std::mem::transmute; @@ -17,6 +17,16 @@ pub struct Resample { ctx: *mut SwrContext, } +impl Drop for Resample { + fn drop(&mut self) { + unsafe { + if !self.ctx.is_null() { + swr_free(&mut self.ctx); + } + } + } +} + impl Resample { pub fn new(format: AVSampleFormat, rate: u32, channels: usize) -> Self { Self { @@ -45,10 +55,10 @@ impl Resample { 0, ptr::null_mut(), ); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); let ret = swr_init(self.ctx); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); Ok(()) } diff --git a/src/scale.rs b/src/scale.rs index d45642c..e0de96d 100644 --- a/src/scale.rs +++ b/src/scale.rs @@ -1,7 +1,7 @@ use std::mem::transmute; use std::ptr; -use crate::{return_ffmpeg_error, rstr}; +use crate::{bail_ffmpeg, rstr}; use anyhow::{bail, Error}; use ffmpeg_sys_the_third::{ av_frame_alloc, av_frame_copy_props, av_get_pix_fmt_name, sws_freeContext, sws_getContext, @@ -16,13 +16,12 @@ pub struct Scaler { ctx: *mut SwsContext, } -unsafe impl Send for Scaler {} - impl Drop for Scaler { fn drop(&mut self) { unsafe { - sws_freeContext(self.ctx); - self.ctx = ptr::null_mut(); + if !self.ctx.is_null() { + sws_freeContext(self.ctx); + } } } } @@ -111,10 +110,10 @@ impl Scaler { let dst_frame = av_frame_alloc(); let ret = av_frame_copy_props(dst_frame, frame); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); let ret = sws_scale_frame(self.ctx, dst_frame, frame); - return_ffmpeg_error!(ret); + bail_ffmpeg!(ret); Ok(dst_frame) } diff --git a/src/transcode.rs b/src/transcode.rs new file mode 100644 index 0000000..2805081 --- /dev/null +++ b/src/transcode.rs @@ -0,0 +1,9 @@ +use crate::{Decoder, Demuxer, Encoder, Muxer, Scaler}; + +pub struct Transcoder { + demuxer: Demuxer, + decoder: Decoder, + scaler: Scaler, + encoder: Encoder, + muxer: Muxer, +}