waycap_rs/encoders/
opus_encoder.rs

1use ffmpeg_next::{self as ffmpeg, Rational};
2use std::collections::VecDeque;
3
4use ringbuf::{
5    traits::{Producer, Split},
6    HeapCons, HeapProd, HeapRb,
7};
8
9use crate::types::audio_frame::EncodedAudioFrame;
10
11use super::audio::{boost_with_rms, AudioEncoder};
12
13pub struct OpusEncoder {
14    encoder: Option<ffmpeg::codec::encoder::Audio>,
15    next_pts: i64,
16    leftover_data: VecDeque<f32>,
17    encoded_samples_recv: Option<HeapCons<EncodedAudioFrame>>,
18    encoded_samples_sender: Option<HeapProd<EncodedAudioFrame>>,
19    capture_timestamps: VecDeque<i64>,
20}
21
22impl OpusEncoder {
23    fn create_encoder() -> crate::types::error::Result<ffmpeg::codec::encoder::Audio> {
24        let encoder_codec = ffmpeg::codec::encoder::find(ffmpeg_next::codec::Id::OPUS)
25            .ok_or(ffmpeg::Error::EncoderNotFound)?;
26
27        let mut encoder_ctx = ffmpeg::codec::context::Context::new_with_codec(encoder_codec)
28            .encoder()
29            .audio()?;
30
31        encoder_ctx.set_rate(48000);
32        encoder_ctx.set_bit_rate(70_000);
33        encoder_ctx.set_format(ffmpeg::format::Sample::F32(
34            ffmpeg_next::format::sample::Type::Packed,
35        ));
36        encoder_ctx.set_time_base(Rational::new(1, 48000));
37        encoder_ctx.set_frame_rate(Some(Rational::new(1, 48000)));
38        encoder_ctx.set_channel_layout(ffmpeg::channel_layout::ChannelLayout::STEREO);
39
40        let mut encoder = encoder_ctx.open()?;
41
42        // Opus frame size is based on n channels so need to update it
43        unsafe {
44            (*encoder.as_mut_ptr()).frame_size =
45                (encoder.frame_size() as i32 * encoder.channels() as i32) as i32;
46        }
47
48        Ok(encoder)
49    }
50}
51
52impl AudioEncoder for OpusEncoder {
53    fn new() -> crate::types::error::Result<Self>
54    where
55        Self: Sized,
56    {
57        let encoder = Self::create_encoder()?;
58        let audio_ring_buffer = HeapRb::<EncodedAudioFrame>::new(5);
59        let (sender, receiver) = audio_ring_buffer.split();
60        Ok(Self {
61            encoder: Some(encoder),
62            next_pts: 0,
63            leftover_data: VecDeque::with_capacity(10),
64            encoded_samples_recv: Some(receiver),
65            encoded_samples_sender: Some(sender),
66            capture_timestamps: VecDeque::with_capacity(10),
67        })
68    }
69
70    fn process(
71        &mut self,
72        mut raw_frame: crate::types::audio_frame::RawAudioFrame,
73    ) -> crate::types::error::Result<()> {
74        if let Some(ref mut encoder) = self.encoder {
75            let n_channels = encoder.channels() as usize;
76            let total_samples = raw_frame.samples.len();
77
78            if total_samples % n_channels != 0 {
79                return Err(crate::types::error::WaycapError::FFmpeg(
80                    ffmpeg::Error::InvalidData,
81                ));
82            }
83
84            let frame_size = encoder.frame_size() as usize;
85
86            // Boost the audio so that even if system audio level is low
87            // it's still audible in playback
88            boost_with_rms(&mut raw_frame.samples)?;
89            self.leftover_data.extend(raw_frame.samples);
90
91            // Send chunked frames to encoder
92            while self.leftover_data.len() >= frame_size {
93                let frame_samples: Vec<f32> = self.leftover_data.drain(..frame_size).collect();
94                let mut frame = ffmpeg::frame::Audio::new(
95                    encoder.format(),
96                    frame_size,
97                    encoder.channel_layout(),
98                );
99
100                // Capture time in vec
101                frame.plane_mut(0).copy_from_slice(&frame_samples);
102                frame.set_pts(Some(self.next_pts));
103                frame.set_rate(encoder.rate());
104
105                self.capture_timestamps.push_back(raw_frame.timestamp);
106                encoder.send_frame(&frame)?;
107
108                // Try and get a frame back from encoder
109                let mut packet = ffmpeg::codec::packet::Packet::empty();
110                if encoder.receive_packet(&mut packet).is_ok() {
111                    if let Some(data) = packet.data() {
112                        let pts = packet.pts().unwrap_or(0);
113                        if let Some(ref mut sender) = self.encoded_samples_sender {
114                            if sender
115                                .try_push(EncodedAudioFrame {
116                                    data: data.to_vec(),
117                                    pts,
118                                    timestamp: self.capture_timestamps.pop_front().unwrap_or(0),
119                                })
120                                .is_err()
121                            {
122                                log::error!("Could not send encoded packet to audio ringbuf");
123                            }
124                        }
125                    }
126                }
127
128                self.next_pts += frame_size as i64;
129            }
130        }
131
132        Ok(())
133    }
134
135    fn get_encoder(&self) -> &Option<ffmpeg_next::codec::encoder::Audio> {
136        &self.encoder
137    }
138
139    fn drain(&mut self) -> crate::types::error::Result<()> {
140        if let Some(ref mut encoder) = self.encoder {
141            encoder.send_eof()?;
142            let mut packet = ffmpeg::codec::packet::Packet::empty();
143            while encoder.receive_packet(&mut packet).is_ok() {
144                if let Some(data) = packet.data() {
145                    let pts = packet.pts().unwrap_or(0);
146                    if let Some(ref mut sender) = self.encoded_samples_sender {
147                        if sender
148                            .try_push(EncodedAudioFrame {
149                                data: data.to_vec(),
150                                pts,
151                                timestamp: self.capture_timestamps.pop_front().unwrap_or(0),
152                            })
153                            .is_err()
154                        {
155                            log::error!("Could not send encoded packet to audio ringbuf");
156                        }
157                    }
158                }
159            }
160        }
161
162        Ok(())
163    }
164
165    fn drop_encoder(&mut self) {
166        self.encoder.take();
167    }
168
169    fn reset(&mut self) -> crate::types::error::Result<()> {
170        self.drop_encoder();
171        self.capture_timestamps.clear();
172        self.encoder = Some(Self::create_encoder()?);
173
174        Ok(())
175    }
176
177    fn take_encoded_recv(&mut self) -> Option<HeapCons<EncodedAudioFrame>> {
178        self.encoded_samples_recv.take()
179    }
180}
waycap_rs/encoders/opus_encoder.rs

waycap_rs/encoders/
opus_encoder.rs
