vidformer/av/
muxer.rs

1use log::*;
2use num_rational::Rational64;
3use rusty_ffmpeg::ffi;
4use std::ffi::CString;
5use std::ptr;
6
7pub struct Muxer {
8    pub ofmt_ctx: *mut ffi::AVFormatContext,
9    pub out_time_base: ffi::AVRational,
10    pub out_stream: *mut ffi::AVStream,
11    pub frames: usize,
12}
13
14impl Muxer {
15    pub fn new(
16        output_path: &str,
17        codecpar: *mut ffi::AVCodecParameters,
18        time_base: &Rational64,
19        format_name: Option<&str>,
20    ) -> Result<Self, crate::Error> {
21        let output_path = CString::new(output_path).unwrap();
22        let mut ofmt_ctx: *mut ffi::AVFormatContext = ptr::null_mut();
23
24        let format_name_cstr: Option<CString> = format_name.map(|name| CString::new(name).unwrap());
25
26        if unsafe {
27            ffi::avformat_alloc_output_context2(
28                &mut ofmt_ctx,
29                ptr::null_mut(),
30                format_name_cstr
31                    .as_ref()
32                    .map_or(ptr::null(), |s| s.as_ptr()),
33                output_path.as_ptr(),
34            )
35        } < 0
36        {
37            return Err(crate::Error::AVError(
38                "Failed to allocate output context".to_string(),
39            ));
40        }
41
42        unsafe {
43            (*ofmt_ctx).debug = 1;
44        }
45
46        let out_stream = unsafe { ffi::avformat_new_stream(ofmt_ctx, ptr::null()) };
47        if out_stream.is_null() {
48            return Err(crate::Error::AVError(
49                "Failed to allocate output stream".to_string(),
50            ));
51        }
52
53        let time_base = crate::util::rat_to_avrat(time_base);
54        unsafe {
55            ffi::avcodec_parameters_free(&mut (*out_stream).codecpar); // free existing codec parameters, this is the only reference to them
56
57            (*out_stream).codecpar = ffi::avcodec_parameters_alloc();
58            if (*out_stream).codecpar.is_null() {
59                return Err(crate::Error::AVError(
60                    "Failed to allocate codec parameters".to_string(),
61                ));
62            }
63            if ffi::avcodec_parameters_copy((*out_stream).codecpar, codecpar) < 0 {
64                return Err(crate::Error::AVError(
65                    "Failed to copy codec parameters".to_string(),
66                ));
67            }
68
69            (*out_stream).time_base = time_base;
70        }
71
72        if unsafe { (*(*ofmt_ctx).oformat).flags } & ffi::AVFMT_NOFILE as i32 == 0
73            && unsafe {
74                ffi::avio_open(
75                    &mut (*ofmt_ctx).pb,
76                    output_path.as_ptr(),
77                    ffi::AVIO_FLAG_WRITE as i32,
78                )
79            } < 0
80        {
81            return Err(crate::Error::AVError(
82                "Failed to open output file".to_string(),
83            ));
84        }
85
86        // show output on terminal
87        unsafe {
88            ffi::av_dump_format(ofmt_ctx, 0, output_path.as_ptr(), 1);
89        }
90
91        if unsafe { ffi::avformat_write_header(ofmt_ctx, ptr::null_mut()) } < 0 {
92            return Err(crate::Error::AVError(
93                "Failed to write output container header".to_string(),
94            ));
95        }
96
97        Ok(Muxer {
98            ofmt_ctx,
99            out_time_base: unsafe { (*out_stream).time_base },
100            out_stream,
101            frames: 0,
102        })
103    }
104
105    pub fn mux_packet(&mut self, packet: *mut ffi::AVPacket) -> Result<(), crate::Error> {
106        let input_pts = unsafe { (*packet).pts };
107        let input_dts = unsafe { (*packet).dts };
108
109        unsafe {
110            (*packet).pos = -1;
111            (*packet).stream_index = (*self.out_stream).index;
112            (*packet).duration = 0;
113        }
114
115        info!(
116            "MUX - Write packet with pts {} and dts {} side_elements={} flags={} (output pts = {}/{} and dts = {}/{})",
117            input_pts,
118            input_dts,
119            unsafe { (*packet).side_data_elems },
120            unsafe { (*packet).flags },
121            unsafe { (*packet).pts },
122            unsafe { (*self.out_stream).time_base.den},
123            unsafe { (*packet).dts },
124            unsafe { (*self.out_stream).time_base.den},
125        );
126
127        // Todo: maybe optionally interleave?
128        if unsafe { ffi::av_interleaved_write_frame(self.ofmt_ctx, packet) } < 0 {
129            return Err(crate::Error::AVError(
130                "Failed to write packet to output".to_string(),
131            ));
132        }
133
134        self.frames += 1;
135
136        Ok(())
137    }
138
139    pub fn close(self) -> Result<(), crate::Error> {
140        info!("Closing muxer");
141
142        // flush
143        if unsafe { ffi::av_interleaved_write_frame(self.ofmt_ctx, ptr::null_mut()) } < 0 {
144            return Err(crate::Error::AVError("Failed to flush output".to_string()));
145        }
146
147        unsafe {
148            ffi::av_write_trailer(self.ofmt_ctx);
149            ffi::avio_close((*self.ofmt_ctx).pb);
150            ffi::avformat_free_context(self.ofmt_ctx);
151        }
152
153        Ok(())
154    }
155}
vidformer/av/muxer.rs

vidformer/av/
muxer.rs
