ffmpeg_decoder/
decoder.rs

1use crate::error::Error;
2
3use ffmpeg_sys_next::{
4    self, av_frame_alloc, av_frame_free, av_frame_unref, av_freep, av_get_alt_sample_fmt,
5    av_get_bytes_per_sample, av_get_channel_layout_nb_channels, av_get_sample_fmt_name,
6    av_init_packet, av_packet_unref, av_read_frame, av_register_all, av_sample_fmt_is_planar,
7    av_samples_alloc, av_samples_get_buffer_size, avcodec_alloc_context3, avcodec_close,
8    avcodec_find_decoder, avcodec_free_context, avcodec_open2, avcodec_parameters_to_context,
9    avcodec_receive_frame, avcodec_send_packet, avformat_close_input, avformat_find_stream_info,
10    avformat_open_input, swr_alloc_set_opts, swr_convert, swr_get_out_samples, swr_init, AVCodec,
11    AVCodecContext, AVFormatContext, AVFrame, AVMediaType, AVPacket, AVSampleFormat, AVStream,
12};
13use std::ffi::{CStr, CString};
14use std::path::Path;
15use std::ptr;
16use std::slice;
17use std::time::Duration;
18
19use log::{error, info};
20
21const AVERROR_EOF: i32 = -0x20_464_F45;
22const AVERROR_EAGAIN: i32 = -11;
23const AVERROR_EDEADLK: i32 = -35;
24const DEFAULT_CONVERSION_FORMAT: AVSampleFormat = AVSampleFormat::AV_SAMPLE_FMT_S16;
25
26pub struct Decoder {
27    format_ctx: FormatContext,
28    stream: Stream,
29    codec_ctx: CodecContext,
30    frame: Frame,
31    packet: Packet,
32    swr_ctx: Option<SwrContext>,
33    current_frame: Vec<u8>,
34    first_frame_stored: bool,
35}
36
37impl Decoder {
38    pub fn open(path: impl AsRef<Path>) -> Result<Decoder, Error> {
39        unsafe { av_register_all() };
40
41        // Open the file and get the format context
42        let format_ctx = FormatContext::open(&path.as_ref().display().to_string())?;
43
44        // Find first audio stream in file
45        format_ctx.find_stream_info()?;
46        let stream = format_ctx.get_audio_stream()?;
47
48        // Get the streams codec
49        let codec = stream.get_codec()?;
50
51        // Setup codec context and intialize
52        let codec_ctx = codec.get_context()?;
53        codec_ctx.copy_parameters_from_stream(&stream)?;
54        codec_ctx.request_non_planar_format();
55        codec_ctx.initialize()?;
56
57        print_codec_info(&codec_ctx);
58
59        // Allocate frame
60        let frame = Frame::new()?;
61
62        // Initialize packet
63        let packet = Packet::new();
64
65        // Initialize swr context, if conversion is needed
66        let swr_ctx = if codec_ctx.sample_format() != DEFAULT_CONVERSION_FORMAT {
67            Some(SwrContext::new(&codec_ctx)?)
68        } else {
69            None
70        };
71
72        Ok(Decoder {
73            format_ctx,
74            stream,
75            codec_ctx,
76            frame,
77            packet,
78            swr_ctx,
79            current_frame: vec![],
80            first_frame_stored: false,
81        })
82    }
83
84    fn read_next_frame(&mut self) -> ReadFrameStatus {
85        let status =
86            unsafe { av_read_frame(self.format_ctx.inner, self.packet.inner.as_mut_ptr()) };
87
88        match status {
89            AVERROR_EOF => ReadFrameStatus::Eof,
90            _ if status != 0 => ReadFrameStatus::Other(status),
91            _ => ReadFrameStatus::Ok,
92        }
93    }
94
95    fn send_packet_for_decoding(&mut self) -> SendPacketStatus {
96        let status =
97            unsafe { avcodec_send_packet(self.codec_ctx.inner, self.packet.inner.as_mut_ptr()) };
98
99        match status {
100            0 => SendPacketStatus::Ok,
101            _ => SendPacketStatus::Other(status),
102        }
103    }
104
105    fn receive_decoded_frame(&self) -> ReceiveFrameStatus {
106        let status = unsafe { avcodec_receive_frame(self.codec_ctx.inner, self.frame.inner) };
107
108        match status {
109            0 => ReceiveFrameStatus::Ok,
110            AVERROR_EAGAIN => ReceiveFrameStatus::Again,
111            AVERROR_EDEADLK => ReceiveFrameStatus::Deadlk,
112            _ => ReceiveFrameStatus::Other(status),
113        }
114    }
115
116    fn convert_and_store_frame(&mut self) {
117        let num_samples = self.frame.num_samples();
118        let channel_layout = self.frame.channel_layout();
119        let num_channels = unsafe { av_get_channel_layout_nb_channels(channel_layout) };
120
121        let extended_data = self.frame.extended_data();
122
123        let mut out_buf = std::ptr::null_mut::<u8>();
124
125        let out_slice = if self.swr_ctx.is_some() {
126            let out_samples =
127                unsafe { swr_get_out_samples(self.swr_ctx.as_ref().unwrap().inner, num_samples) };
128
129            unsafe {
130                av_samples_alloc(
131                    &mut out_buf,
132                    ptr::null_mut(),
133                    num_channels,
134                    out_samples,
135                    DEFAULT_CONVERSION_FORMAT,
136                    0,
137                )
138            };
139
140            unsafe {
141                swr_convert(
142                    self.swr_ctx.as_ref().unwrap().inner,
143                    &mut out_buf,
144                    out_samples,
145                    extended_data,
146                    num_samples,
147                )
148            };
149
150            let out_size = unsafe {
151                av_samples_get_buffer_size(
152                    ptr::null_mut(),
153                    num_channels,
154                    out_samples,
155                    DEFAULT_CONVERSION_FORMAT,
156                    0,
157                )
158            };
159
160            unsafe { slice::from_raw_parts(out_buf, out_size as usize) }
161        } else {
162            unsafe {
163                slice::from_raw_parts(
164                    extended_data.as_ref().unwrap().as_ref().unwrap(),
165                    self.frame.inner.as_ref().unwrap().linesize[0] as usize,
166                )
167            }
168        };
169
170        if !self.current_frame.is_empty() {
171            self.current_frame.drain(..);
172        }
173
174        self.current_frame.extend_from_slice(out_slice);
175
176        if self.swr_ctx.is_some() {
177            // Free samples buffer
178            unsafe { av_freep(&mut out_buf as *mut _ as _) };
179        }
180
181        unsafe { av_frame_unref(self.frame.inner) };
182    }
183
184    fn frame_for_stream(&self) -> bool {
185        unsafe { self.packet.inner.as_ptr().as_ref().unwrap().stream_index == self.stream.index }
186    }
187
188    fn reset_packet(&mut self) {
189        unsafe { av_packet_unref(self.packet.inner.as_mut_ptr()) };
190    }
191
192    fn next_sample(&mut self) -> i16 {
193        let sample_u8: [u8; 2] = [self.current_frame.remove(0), self.current_frame.remove(0)];
194
195        ((sample_u8[1] as i16) << 8) | sample_u8[0] as i16
196    }
197
198    fn process_next_frame(&mut self) -> Option<Result<(), Error>> {
199        match self.read_next_frame() {
200            ReadFrameStatus::Ok => {}
201            ReadFrameStatus::Eof => {
202                return None;
203            }
204            ReadFrameStatus::Other(status) => {
205                error!("{}", Error::ReadFrame(status));
206                return None;
207            }
208        }
209
210        if !self.frame_for_stream() {
211            self.reset_packet();
212            return self.process_next_frame();
213        }
214
215        match self.send_packet_for_decoding() {
216            SendPacketStatus::Ok => self.reset_packet(),
217            SendPacketStatus::Other(status) => {
218                error!("{}", Error::SendPacket(status));
219                return None;
220            }
221        }
222
223        match self.receive_decoded_frame() {
224            ReceiveFrameStatus::Ok => {}
225            ReceiveFrameStatus::Again | ReceiveFrameStatus::Deadlk => {
226                return self.process_next_frame()
227            }
228            ReceiveFrameStatus::Other(status) => {
229                error!("{}", Error::ReceiveFrame(status));
230                return None;
231            }
232        }
233
234        self.convert_and_store_frame();
235
236        Some(Ok(()))
237    }
238
239    fn cleanup(&mut self) {
240        // Drain the decoder.
241        drain_decoder(self.codec_ctx.inner).unwrap();
242
243        unsafe {
244            // Free all data used by the frame.
245            av_frame_free(&mut self.frame.inner);
246
247            // Close the context and free all data associated to it, but not the context itself.
248            avcodec_close(self.codec_ctx.inner);
249
250            // Free the context itself.
251            avcodec_free_context(&mut self.codec_ctx.inner);
252
253            // Close the input.
254            avformat_close_input(&mut self.format_ctx.inner);
255        }
256    }
257
258    pub(crate) fn _current_frame_len(&self) -> Option<usize> {
259        Some(self.current_frame.len())
260    }
261
262    pub(crate) fn _channels(&self) -> u16 {
263        self.codec_ctx.channels() as _
264    }
265
266    pub(crate) fn _sample_rate(&self) -> u32 {
267        self.codec_ctx.sample_rate() as _
268    }
269
270    pub(crate) fn _total_duration(&self) -> Option<Duration> {
271        //TODO let duration = self.stream.duration();
272        None
273    }
274}
275
276unsafe impl Send for Decoder {}
277
278impl Iterator for Decoder {
279    type Item = i16;
280
281    #[inline]
282    fn next(&mut self) -> Option<Self::Item> {
283        if !self.first_frame_stored {
284            if self.process_next_frame().is_none() {
285                self.cleanup();
286                return None;
287            }
288
289            self.first_frame_stored = true;
290
291            return Some(self.next_sample());
292        }
293
294        if !self.current_frame.is_empty() {
295            return Some(self.next_sample());
296        }
297
298        match self.receive_decoded_frame() {
299            ReceiveFrameStatus::Ok => {
300                self.convert_and_store_frame();
301                Some(self.next_sample())
302            }
303            ReceiveFrameStatus::Again | ReceiveFrameStatus::Deadlk => {
304                if self.process_next_frame().is_none() {
305                    self.cleanup();
306                    return None;
307                }
308
309                Some(self.next_sample())
310            }
311            ReceiveFrameStatus::Other(status) => {
312                error!("{}", Error::ReceiveFrame(status));
313                self.cleanup();
314                None
315            }
316        }
317    }
318}
319
320struct FormatContext {
321    inner: *mut AVFormatContext,
322}
323
324impl FormatContext {
325    fn open(path: &str) -> Result<FormatContext, Error> {
326        let mut inner = std::ptr::null_mut::<AVFormatContext>();
327
328        let path = CString::new(path).unwrap();
329
330        let status = unsafe {
331            avformat_open_input(
332                &mut inner,
333                path.as_ptr(),
334                std::ptr::null_mut(),
335                std::ptr::null_mut(),
336            )
337        };
338        if status != 0 {
339            return Err(Error::InitializeFormatContext);
340        }
341
342        Ok(FormatContext { inner })
343    }
344
345    /// Look at first few frames to determine stream info
346    fn find_stream_info(&self) -> Result<(), Error> {
347        let status = unsafe { avformat_find_stream_info(self.inner, ptr::null_mut()) };
348        if status < 0 {
349            return Err(Error::FindStreamInfo);
350        }
351        Ok(())
352    }
353
354    ///  Get the first audio stream
355    fn get_audio_stream(&self) -> Result<Stream, Error> {
356        let num_streams = unsafe { self.inner.as_ref().unwrap().nb_streams };
357        let streams = unsafe { self.inner.as_ref().unwrap().streams };
358
359        let streams = unsafe { slice::from_raw_parts(streams, num_streams as usize) };
360
361        let stream_idx = find_audio_stream(streams)?;
362
363        Ok(Stream::new(streams[0], stream_idx))
364    }
365}
366
367struct SwrContext {
368    inner: *mut ffmpeg_sys_next::SwrContext,
369}
370
371impl SwrContext {
372    fn new(codec_ctx: &CodecContext) -> Result<SwrContext, Error> {
373        let swr_ctx: *mut ffmpeg_sys_next::SwrContext = unsafe {
374            swr_alloc_set_opts(
375                ptr::null_mut(),
376                codec_ctx.channel_layout() as i64,
377                DEFAULT_CONVERSION_FORMAT,
378                codec_ctx.sample_rate(),
379                codec_ctx.channel_layout() as i64,
380                codec_ctx.sample_format(),
381                codec_ctx.sample_rate(),
382                0,
383                ptr::null_mut(),
384            )
385        };
386
387        let status = unsafe { swr_init(swr_ctx) };
388        if status != 0 {
389            return Err(Error::InitializeSwr);
390        }
391
392        Ok(SwrContext { inner: swr_ctx })
393    }
394}
395
396struct Packet {
397    inner: std::mem::MaybeUninit<AVPacket>,
398}
399
400impl Packet {
401    fn new() -> Packet {
402        let mut packet = std::mem::MaybeUninit::uninit();
403
404        unsafe { av_init_packet(packet.as_mut_ptr()) };
405
406        Packet { inner: packet }
407    }
408}
409
410struct Frame {
411    inner: *mut AVFrame,
412}
413
414impl Frame {
415    fn new() -> Result<Frame, Error> {
416        let frame: *mut AVFrame = unsafe { av_frame_alloc() };
417
418        if frame.is_null() {
419            return Err(Error::NullFrame);
420        }
421
422        Ok(Frame { inner: frame })
423    }
424
425    fn num_samples(&self) -> i32 {
426        unsafe { self.inner.as_ref().unwrap().nb_samples }
427    }
428
429    fn channel_layout(&self) -> u64 {
430        unsafe { self.inner.as_ref().unwrap().channel_layout }
431    }
432
433    fn extended_data(&self) -> *mut *const u8 {
434        unsafe { self.inner.as_ref().unwrap().extended_data as *mut *const u8 }
435    }
436}
437
438struct Stream {
439    inner: *mut AVStream,
440    index: i32,
441}
442
443impl Stream {
444    fn new(inner: *mut AVStream, index: i32) -> Stream {
445        Stream { inner, index }
446    }
447
448    fn get_codec(&self) -> Result<Codec, Error> {
449        // Get streams codec
450        let codec_params = unsafe { self.inner.as_ref().unwrap().codecpar };
451        let codec_id = unsafe { codec_params.as_ref().unwrap().codec_id };
452
453        let codec: *mut AVCodec = unsafe { avcodec_find_decoder(codec_id) };
454        if codec.is_null() {
455            return Err(Error::NullCodec);
456        }
457
458        Ok(Codec::new(codec))
459    }
460
461    #[allow(dead_code)]
462    fn duration(&self) -> i64 {
463        unsafe { self.inner.as_ref().unwrap().duration }
464    }
465}
466
467struct CodecContext {
468    inner: *mut AVCodecContext,
469    codec: *mut AVCodec,
470}
471
472impl CodecContext {
473    fn new(inner: *mut AVCodecContext, codec: *mut AVCodec) -> CodecContext {
474        CodecContext { inner, codec }
475    }
476
477    fn copy_parameters_from_stream(&self, stream: &Stream) -> Result<(), Error> {
478        let params = unsafe { stream.inner.as_ref().unwrap().codecpar };
479
480        let status = unsafe { avcodec_parameters_to_context(self.inner, params) };
481
482        if status != 0 {
483            return Err(Error::CodecParamsToContext);
484        }
485
486        Ok(())
487    }
488
489    fn request_non_planar_format(&self) {
490        unsafe {
491            let sample_fmt = self.inner.as_ref().unwrap().sample_fmt;
492            let alt_format = av_get_alt_sample_fmt(sample_fmt, 0);
493
494            self.inner.as_mut().unwrap().request_sample_fmt = alt_format;
495        }
496    }
497
498    fn initialize(&self) -> Result<(), Error> {
499        let status = unsafe { avcodec_open2(self.inner, self.codec, &mut std::ptr::null_mut()) };
500
501        if status != 0 {
502            return Err(Error::InitializeDecoder);
503        }
504
505        Ok(())
506    }
507
508    fn codec_name(&self) -> &str {
509        let name = unsafe { CStr::from_ptr(self.codec.as_ref().unwrap().long_name) };
510
511        name.to_str().unwrap()
512    }
513
514    fn sample_format(&self) -> AVSampleFormat {
515        unsafe { self.inner.as_ref().unwrap().sample_fmt }
516    }
517
518    fn sample_format_name(&self) -> &str {
519        let sample_fmt = unsafe { CStr::from_ptr(av_get_sample_fmt_name(self.sample_format())) };
520
521        sample_fmt.to_str().unwrap()
522    }
523
524    fn sample_rate(&self) -> i32 {
525        unsafe { self.inner.as_ref().unwrap().sample_rate }
526    }
527
528    fn sample_size(&self) -> i32 {
529        unsafe { av_get_bytes_per_sample(self.inner.as_ref().unwrap().sample_fmt) }
530    }
531
532    fn channels(&self) -> i32 {
533        unsafe { self.inner.as_ref().unwrap().channels }
534    }
535
536    fn channel_layout(&self) -> u64 {
537        unsafe { self.inner.as_ref().unwrap().channel_layout }
538    }
539
540    fn is_planar(&self) -> i32 {
541        unsafe { av_sample_fmt_is_planar(self.inner.as_ref().unwrap().sample_fmt) }
542    }
543}
544
545struct Codec {
546    inner: *mut AVCodec,
547}
548
549impl Codec {
550    fn new(inner: *mut AVCodec) -> Codec {
551        Codec { inner }
552    }
553
554    fn get_context(&self) -> Result<CodecContext, Error> {
555        let ctx: *mut AVCodecContext = unsafe { avcodec_alloc_context3(self.inner) };
556
557        if ctx.is_null() {
558            return Err(Error::NullCodecContext);
559        }
560
561        Ok(CodecContext::new(ctx, self.inner))
562    }
563}
564
565enum ReadFrameStatus {
566    Ok,
567    Eof,
568    Other(i32),
569}
570
571enum SendPacketStatus {
572    Ok,
573    Other(i32),
574}
575
576enum ReceiveFrameStatus {
577    Ok,
578    Again,
579    Deadlk,
580    Other(i32),
581}
582
583fn find_audio_stream(streams: &[*mut AVStream]) -> Result<i32, Error> {
584    for stream in streams {
585        let codec_type = unsafe {
586            stream
587                .as_ref()
588                .unwrap()
589                .codecpar
590                .as_ref()
591                .unwrap()
592                .codec_type
593        };
594        let index = unsafe { stream.as_ref().unwrap().index };
595
596        if codec_type == AVMediaType::AVMEDIA_TYPE_AUDIO {
597            return Ok(index);
598        }
599    }
600
601    Err(Error::NoAudioStream)
602}
603
604fn print_codec_info(codec_ctx: &CodecContext) {
605    info!("Codec:         {}", codec_ctx.codec_name());
606    info!("Sample Format: {}", codec_ctx.sample_format_name());
607    info!("Sample Rate:   {}", codec_ctx.sample_rate());
608    info!("Sample Size:   {}", codec_ctx.sample_size());
609    info!("Channels:      {}", codec_ctx.channels());
610    info!("Planar:        {}", codec_ctx.is_planar());
611}
612
613fn drain_decoder(codec_ctx: *mut AVCodecContext) -> Result<(), Error> {
614    let status = unsafe { avcodec_send_packet(codec_ctx, std::ptr::null()) };
615    if status == 0 {
616    } else {
617        return Err(Error::DrainDecoder(status));
618    }
619
620    Ok(())
621}
ffmpeg_decoder/decoder.rs

ffmpeg_decoder/
decoder.rs
