waycap_rs/encoders/
opus_encoder.rs

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

waycap_rs/encoders/
opus_encoder.rs
