1use std::collections::VecDeque;
2use std::sync::Arc;
3
4#[cfg(feature = "audio")]
5use crate::audio::{MpaAudioChunk, MpaAudioDecoder};
6use crate::convert::frame_to_rgba_bt601_limited;
7use crate::demux::{Demuxer, Packet, StreamType};
8use crate::error::Result;
9use crate::video::{Decoder as VideoDecoder, Frame};
10
11#[derive(Clone)]
12pub struct MpegRgbaFrame {
13 pub pts_ms: i64,
14 pub width: u32,
15 pub height: u32,
16 pub rgba: Vec<u8>,
17}
18
19#[derive(Clone)]
23pub struct MpegAudioF32 {
24 pub pts_ms: i64,
25 pub sample_rate: u32,
26 pub channels: u16,
27 pub samples: Vec<f32>,
28}
29
30#[derive(Clone)]
31pub enum MpegAvEvent {
32 Video(MpegRgbaFrame),
33 Audio(MpegAudioF32),
34}
35
36#[derive(Default)]
37pub struct MpegAvPipeline {
38 demux: Demuxer,
39 vdec: VideoDecoder,
40 #[cfg(feature = "audio")]
41 adec: MpaAudioDecoder,
42
43 pkts: Vec<Packet>,
44 pub stash: VecDeque<MpegAvEvent>,
45}
46
47impl MpegAvPipeline {
48 pub fn new() -> Self {
49 Self {
50 demux: Demuxer::new_auto(),
51 vdec: VideoDecoder::new(),
52 #[cfg(feature = "audio")]
53 adec: MpaAudioDecoder::new(),
54 pkts: Vec::new(),
55 stash: VecDeque::new(),
56 }
57 }
58
59 #[inline]
60 pub fn demuxer_mut(&mut self) -> &mut Demuxer {
61 &mut self.demux
62 }
63
64 #[inline]
65 pub fn video_decoder_mut(&mut self) -> &mut VideoDecoder {
66 &mut self.vdec
67 }
68
69 #[cfg(feature = "audio")]
70 #[inline]
71 pub fn audio_decoder_mut(&mut self) -> &mut MpaAudioDecoder {
72 &mut self.adec
73 }
74
75 pub fn push_with<F>(&mut self, data: &[u8], pts_90k: Option<i64>, mut on_event: F) -> Result<()>
76 where
77 F: FnMut(MpegAvEvent),
78 {
79 self.pkts.clear();
80 self.demux.push_into(data, pts_90k, &mut self.pkts);
81
82 let mut local_pkts: Vec<Packet> = Vec::new();
84 std::mem::swap(&mut self.pkts, &mut local_pkts);
85
86 for pkt in local_pkts.drain(..) {
87 match pkt.stream_type {
88 StreamType::MpegVideo => self.handle_video_pkt(&pkt, &mut on_event)?,
89 StreamType::MpegAudio => {
90 #[cfg(feature = "audio")]
91 self.handle_audio_pkt(&pkt, &mut on_event)?;
92 }
93 StreamType::Unknown => {}
94 }
95 }
96
97 std::mem::swap(&mut self.pkts, &mut local_pkts);
98 self.pkts.clear();
99
100 Ok(())
101 }
102
103 pub fn push(&mut self, data: &[u8], pts_90k: Option<i64>) -> Result<()> {
104 let mut tmp: Vec<MpegAvEvent> = Vec::new();
105 self.push_with(data, pts_90k, |ev| tmp.push(ev))?;
106 for ev in tmp {
107 self.stash.push_back(ev);
108 }
109 Ok(())
110 }
111
112 pub fn flush_with<F>(&mut self, mut on_event: F) -> Result<()>
113 where
114 F: FnMut(MpegAvEvent),
115 {
116 for f in self.vdec.flush_shared()? {
118 self.emit_video_frame(f, &mut on_event)?;
119 }
120 Ok(())
121 }
122
123 pub fn flush(&mut self) -> Result<()> {
124 let mut tmp: Vec<MpegAvEvent> = Vec::new();
125 self.flush_with(|ev| tmp.push(ev))?;
126 for ev in tmp {
127 self.stash.push_back(ev);
128 }
129 Ok(())
130 }
131
132 fn handle_video_pkt<F>(&mut self, pkt: &Packet, on_event: &mut F) -> Result<()>
133 where
134 F: FnMut(MpegAvEvent),
135 {
136 let decoded: Vec<Arc<Frame>> = self.vdec.decode_shared(&pkt.data, pkt.pts_90k)?;
137 for f in decoded {
138 self.emit_video_frame(f, on_event)?;
139 }
140 Ok(())
141 }
142
143 fn emit_video_frame<F>(&mut self, f: Arc<Frame>, on_event: &mut F) -> Result<()>
144 where
145 F: FnMut(MpegAvEvent),
146 {
147 let w = f.width as u32;
148 let h = f.height as u32;
149 let mut rgba = vec![0u8; (w as usize) * (h as usize) * 4];
150 frame_to_rgba_bt601_limited(&f, &mut rgba);
151
152 let pts_ms = pts90k_opt_to_ms(f.pts_90k);
153 on_event(MpegAvEvent::Video(MpegRgbaFrame {
154 pts_ms,
155 width: w,
156 height: h,
157 rgba,
158 }));
159 Ok(())
160 }
161
162 #[cfg(feature = "audio")]
163 fn handle_audio_pkt<F>(&mut self, pkt: &Packet, on_event: &mut F) -> Result<()>
164 where
165 F: FnMut(MpegAvEvent),
166 {
167 let pts_ms_opt = pkt.pts_90k.map(pts90k_to_ms);
168 self.adec
169 .push_with(&pkt.data, pts_ms_opt, |ch: MpaAudioChunk| {
170 on_event(MpegAvEvent::Audio(MpegAudioF32 {
171 pts_ms: ch.pts_ms,
172 sample_rate: ch.sample_rate,
173 channels: ch.channels,
174 samples: ch.samples,
175 }))
176 })?;
177 Ok(())
178 }
179}
180
181#[inline]
182fn pts90k_to_ms(v: i64) -> i64 {
183 (v * 1000) / 90000
184}
185
186#[inline]
187fn pts90k_opt_to_ms(v: Option<i64>) -> i64 {
188 v.map(pts90k_to_ms).unwrap_or(0)
189}