feat: muxer

.gitignore

@@ -1,4 +1,7 @@
 /target
 /.idea
 
-test.png
+# Test output
+*.mp4
+*.png
+*.mkv

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: {} <path>", name);

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);
+            if !self.context.is_null() {
-            self.codec = ptr::null_mut();
+                avcodec_free_context(&mut self.context);
-            self.context = ptr::null_mut();
+            }
         }
     }
 }
@@ -66,8 +66,6 @@ impl Display for DecoderCodecContext {
     }
 }
 
-unsafe impl Send for DecoderCodecContext {}
-
 pub struct Decoder {
     /// Decoder instances by stream index
     codecs: HashMap<i32, DecoderCodecContext>,
@@ -75,6 +73,12 @@ pub struct Decoder {
     hw_decoder_types: Option<HashSet<AVHWDeviceType>>,
 }
 
+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<Vec<(*mut AVFrame, *mut AVStream)>, 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 {

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<DemuxerInfo, Error> {
         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"));

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,
+    av_channel_layout_default, av_d2q, av_inv_q, av_packet_alloc, av_packet_free,
-    avcodec_find_encoder, avcodec_get_supported_config, avcodec_open2, avcodec_receive_packet,
+    avcodec_alloc_context3, avcodec_find_encoder, avcodec_free_context,
-    avcodec_send_frame, AVChannelLayout, AVCodec, AVCodecConfig, AVCodecContext, AVCodecID,
+    avcodec_get_supported_config, avcodec_open2, avcodec_receive_packet, avcodec_send_frame,
-    AVFrame, AVPacket, AVPixelFormat, AVRational, AVSampleFormat, AVERROR, AVERROR_EOF,
+    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_*)
+    /// Set the encoder framerate
-    pub unsafe fn with_framerate(self, rate: AVRational) -> Self {
+    pub unsafe fn with_framerate(self, fps: f32) -> Self {
-        (*self.ctx).framerate = rate;
+        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<HashMap<String, String>>) -> Result<Self, Error> {
+        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_width((*frame).width)
-                .with_height(512);
+                .with_height((*frame).height);
 
             let pix_fmts: &[AVPixelFormat] =
                 encoder.list_configs(AVCodecConfig::AV_CODEC_CONFIG_PIX_FORMAT)?;

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(())
     }
 

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<String, String>) -> 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<String, String>) -> 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).width = 1024;
-    (*frame).height = 512;
+    (*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::*;

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<HashMap<String, String>>,
+    ) -> Result<Self> {
+        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<Self> {
+        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<Self> {
+        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(())
+    }
+}

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,
+    swr_alloc_set_opts2, swr_convert_frame, swr_free, swr_init, AVChannelLayout, AVFrame,
-    SwrContext,
+    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(())
     }

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);
+            if !self.ctx.is_null() {
-            self.ctx = ptr::null_mut();
+                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)
     }

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,
+}