media_codec_opus/
encoder.rs

1use std::{collections::VecDeque, os::raw::c_int, sync::Arc};
2
3use bytemuck;
4use ctor::ctor;
5use media_codec::{
6    codec::{AudioParameters, Codec, CodecBuilder, CodecID},
7    encoder::{register_encoder, AudioEncoder, AudioEncoderParameters, Encoder, EncoderBuilder, EncoderParameters},
8    packet::Packet,
9    CodecInformation, CodecParameters,
10};
11use media_core::{
12    audio::SampleFormat,
13    buffer::BufferPool,
14    error::Error,
15    frame::{Frame, SharedFrame},
16    invalid_param_error,
17    rational::Rational64,
18    unsupported_error,
19    variant::Variant,
20    Result,
21};
22
23use crate::{opus_error_string, opus_sys};
24
25struct OpusOptions {
26    application: i32,
27    frame_duration: f32,
28    frame_size: u32,
29    packet_loss: i32,
30    fec: bool,
31    vbr: u32,
32    max_bandwidth: u32,
33    complexity: u32,
34}
35
36impl Default for OpusOptions {
37    fn default() -> Self {
38        OpusOptions {
39            application: opus_sys::OPUS_APPLICATION_AUDIO,
40            frame_duration: 20.0,
41            frame_size: 960,
42            packet_loss: 0,
43            fec: false,
44            vbr: 1,
45            max_bandwidth: 0,
46            complexity: 10,
47        }
48    }
49}
50
51impl OpusOptions {
52    fn from_variant(variant: Option<&Variant>) -> Self {
53        if let Some(variant) = variant {
54            let application = variant["application"].get_int32().unwrap_or(opus_sys::OPUS_APPLICATION_AUDIO);
55            let frame_duration = variant["frame_duration"].get_float().unwrap_or(20.0);
56            let packet_loss = variant["packet_loss"].get_int32().unwrap_or(0);
57            let fec = variant["fec"].get_bool().unwrap_or(false);
58            let vbr = variant["vbr"].get_uint32().unwrap_or(1);
59
60            OpusOptions {
61                application,
62                frame_duration,
63                frame_size: (frame_duration * 48000f32 / 1000f32) as u32,
64                packet_loss,
65                fec,
66                vbr,
67                max_bandwidth: 0,
68                complexity: 10,
69            }
70        } else {
71            Self::default()
72        }
73    }
74}
75
76struct OpusEncoder {
77    encoder: *mut opus_sys::OpusEncoder,
78    pending: VecDeque<Packet<'static>>,
79    options: OpusOptions,
80    buffer: Vec<u8>,
81}
82
83unsafe impl Send for OpusEncoder {}
84unsafe impl Sync for OpusEncoder {}
85
86impl Codec<AudioEncoder> for OpusEncoder {
87    fn configure(&mut self, _params: Option<&CodecParameters>, _options: Option<&Variant>) -> Result<()> {
88        Ok(())
89    }
90
91    fn set_option(&mut self, key: &str, value: &Variant) -> Result<()> {
92        let value = value.get_int32().ok_or_else(|| invalid_param_error!(value))?;
93
94        match key {
95            "bit_rate" => self.encoder_ctl(opus_sys::OPUS_SET_BITRATE_REQUEST, value),
96            _ => Err(unsupported_error!(key)),
97        }
98    }
99}
100
101const DEFAULT_PACKET_PENDING_CAPACITY: usize = 8;
102
103// The maximum frame size is 1275 bytes
104const MAX_FRAME_SIZE: usize = 1275;
105// 120ms packets consist of 6 frames in one packet
106const MAX_FRAMES: usize = 6;
107// The packet header size is 7 bytes
108const PACKET_HEADER_SIZE: usize = 7;
109
110impl Encoder<AudioEncoder> for OpusEncoder {
111    fn send_frame(&mut self, _config: &AudioEncoder, pool: Option<&Arc<BufferPool>>, frame: SharedFrame<Frame<'static>>) -> Result<()> {
112        self.encode(frame, pool)?;
113        Ok(())
114    }
115
116    fn receive_packet(&mut self, _parameters: &AudioEncoder, _pool: Option<&Arc<BufferPool>>) -> Result<Packet<'static>> {
117        self.pending.pop_front().ok_or_else(|| Error::Again("no packet available".to_string()))
118    }
119
120    fn receive_packet_borrowed(&mut self, _config: &AudioEncoder) -> Result<Packet<'_>> {
121        Err(Error::Unsupported("borrowed packet".to_string()))
122    }
123
124    fn flush(&mut self, _config: &AudioEncoder) -> Result<()> {
125        self.buffer.fill(0);
126
127        Ok(())
128    }
129}
130
131impl Drop for OpusEncoder {
132    fn drop(&mut self) {
133        unsafe { opus_sys::opus_encoder_destroy(self.encoder) }
134    }
135}
136
137impl OpusEncoder {
138    pub fn new(codec_id: CodecID, parameters: &AudioEncoderParameters, options: Option<&Variant>) -> Result<Self> {
139        if codec_id != CodecID::OPUS {
140            return Err(unsupported_error!(codec_id));
141        }
142
143        let mut opts = OpusOptions::from_variant(options);
144
145        let audio_params = &parameters.audio;
146        let sample_format = audio_params.format.ok_or_else(|| invalid_param_error!(parameters))?;
147
148        if sample_format != SampleFormat::S16 && sample_format != SampleFormat::F32 {
149            return Err(unsupported_error!(sample_format));
150        }
151
152        let sample_rate = audio_params.sample_rate.ok_or_else(|| invalid_param_error!(parameters))?.get() as opus_sys::opus_int32;
153        let channels = audio_params.channel_layout.as_ref().ok_or_else(|| invalid_param_error!(parameters))?.channels.get() as c_int;
154
155        // Calculate frame size in samples at 48kHz to validate frame duration
156        let frame_size = (opts.frame_duration * 48000f32 / 1000f32) as u32;
157        match frame_size {
158            // 2.5ms | 5ms
159            120 | 240 => {
160                if opts.application != opus_sys::OPUS_APPLICATION_RESTRICTED_LOWDELAY {
161                    opts.application = opus_sys::OPUS_APPLICATION_RESTRICTED_LOWDELAY;
162                }
163            }
164            // 10ms | 20ms | 40ms | 60ms | 80ms | 100ms | 120ms
165            480 | 960 | 1920 | 2880 | 3840 | 4800 | 5760 => {}
166            _ => return Err(Error::Invalid("frame duration".into())),
167        }
168
169        opts.frame_size = frame_size * sample_rate as u32 / 48000;
170
171        let mut error = 0;
172        let opus_encoder = unsafe { opus_sys::opus_encoder_create(sample_rate, channels, opts.application, &mut error) };
173        if opus_encoder.is_null() || error != opus_sys::OPUS_OK {
174            return Err(Error::CreationFailed(opus_error_string(error)));
175        }
176
177        let mut encoder: OpusEncoder = OpusEncoder {
178            encoder: opus_encoder,
179            pending: VecDeque::with_capacity(DEFAULT_PACKET_PENDING_CAPACITY),
180            options: opts,
181            buffer: vec![0u8; frame_size as usize * channels as usize * sample_format.bytes() as usize],
182        };
183
184        encoder.set_audio_parameters(audio_params)?;
185        encoder.set_encoder_parameters(&parameters.encoder)?;
186        encoder.update_options()?;
187
188        Ok(encoder)
189    }
190
191    fn encoder_ctl(&mut self, key: i32, value: i32) -> Result<()> {
192        let ret = unsafe { opus_sys::opus_encoder_ctl(self.encoder, key, value) };
193
194        if ret != opus_sys::OPUS_OK {
195            return Err(Error::SetFailed(opus_error_string(ret)));
196        }
197
198        Ok(())
199    }
200
201    fn set_audio_parameters(&mut self, _audio_params: &AudioParameters) -> Result<()> {
202        Ok(())
203    }
204
205    fn set_encoder_parameters(&mut self, encoder_params: &EncoderParameters) -> Result<()> {
206        if let Some(bit_rate) = encoder_params.bit_rate {
207            self.encoder_ctl(opus_sys::OPUS_SET_BITRATE_REQUEST, bit_rate as i32)?;
208        }
209
210        if let Some(level) = encoder_params.level {
211            self.options.complexity = if !(0..=10).contains(&level) {
212                10
213            } else {
214                level as u32
215            };
216        }
217
218        Ok(())
219    }
220
221    fn update_options(&mut self) -> Result<()> {
222        self.encoder_ctl(opus_sys::OPUS_SET_VBR_REQUEST, (self.options.vbr > 0) as i32)?;
223        self.encoder_ctl(opus_sys::OPUS_SET_VBR_CONSTRAINT_REQUEST, (self.options.vbr == 2) as i32)?;
224        self.encoder_ctl(opus_sys::OPUS_SET_PACKET_LOSS_PERC_REQUEST, self.options.packet_loss)?;
225        self.encoder_ctl(opus_sys::OPUS_SET_INBAND_FEC_REQUEST, self.options.fec as i32)?;
226
227        if self.options.complexity > 0 {
228            self.encoder_ctl(opus_sys::OPUS_SET_COMPLEXITY_REQUEST, self.options.complexity as i32)?;
229        }
230
231        if self.options.max_bandwidth > 0 {
232            self.encoder_ctl(opus_sys::OPUS_SET_MAX_BANDWIDTH_REQUEST, self.options.max_bandwidth as i32)?;
233        }
234
235        Ok(())
236    }
237
238    fn encode(&mut self, frame: SharedFrame<Frame<'static>>, pool: Option<&Arc<BufferPool>>) -> Result<()> {
239        let frame = frame.read();
240        let desc = frame.audio_descriptor().ok_or_else(|| Error::Unsupported("media type".to_string()))?;
241        let sample_format = desc.format;
242
243        if sample_format != SampleFormat::S16 && sample_format != SampleFormat::F32 {
244            return Err(unsupported_error!(sample_format));
245        }
246
247        let guard = frame.map().map_err(|_| Error::Invalid("not readable".into()))?;
248        let planes = guard.planes().unwrap();
249        let packet_size = PACKET_HEADER_SIZE + MAX_FRAME_SIZE * MAX_FRAMES;
250        let channels = desc.channels().get() as usize;
251        let sample_size = channels * sample_format.bytes() as usize;
252        let frame_data = planes.plane_data(0).unwrap();
253        let frame_data_size = desc.samples.get() as usize * sample_size;
254        let chunk_size = self.options.frame_size as usize * sample_size;
255        let mut pts = frame.pts.unwrap_or(0);
256
257        self.buffer.fill(0);
258
259        for chunk in frame_data[..frame_data_size].chunks(chunk_size) {
260            let mut packet = if let Some(pool) = pool {
261                Packet::from_buffer(pool.get_buffer_with_length(packet_size))
262            } else {
263                Packet::new(packet_size)
264            };
265
266            let packet_data = packet.data_mut().ok_or_else(|| Error::Invalid("packet not writable".into()))?;
267            let data = if chunk.len() < chunk_size {
268                self.buffer[..chunk.len()].copy_from_slice(chunk);
269                self.buffer.as_slice()
270            } else {
271                chunk
272            };
273
274            let ret = match desc.format {
275                SampleFormat::S16 => {
276                    let data = bytemuck::cast_slice::<u8, i16>(data);
277                    unsafe {
278                        opus_sys::opus_encode(
279                            self.encoder,
280                            data.as_ptr(),
281                            (data.len() / channels) as i32,
282                            packet_data.as_mut_ptr(),
283                            packet_data.len() as i32,
284                        )
285                    }
286                }
287                SampleFormat::F32 => {
288                    let data = bytemuck::cast_slice::<u8, f32>(data);
289                    unsafe {
290                        opus_sys::opus_encode_float(
291                            self.encoder,
292                            data.as_ptr(),
293                            (data.len() / channels) as i32,
294                            packet_data.as_mut_ptr(),
295                            packet_data.len() as i32,
296                        )
297                    }
298                }
299                _ => return Err(unsupported_error!(sample_format)),
300            };
301
302            if ret < 0 {
303                return Err(Error::Failed(opus_error_string(ret)));
304            }
305
306            let samples = self.options.frame_size as i64;
307            let (duration, time_base) = if let Some(time_base) = frame.time_base {
308                let duration = (Rational64::new(samples, desc.sample_rate.get() as i64) / time_base).to_integer();
309                (duration, time_base)
310            } else {
311                let time_base = Rational64::new(1, desc.sample_rate.get() as i64);
312                let duration = samples;
313                (duration, time_base)
314            };
315
316            packet.pts = Some(pts);
317            packet.duration = Some(duration);
318            packet.time_base = Some(time_base);
319            pts += duration;
320
321            packet.truncate(ret as usize)?;
322
323            self.pending.push_back(packet);
324        }
325
326        Ok(())
327    }
328}
329
330const CODEC_NAME: &str = "opus-enc";
331
332pub struct OpusEncoderBuilder;
333
334impl EncoderBuilder<AudioEncoder> for OpusEncoderBuilder {
335    fn new_encoder(&self, codec_id: CodecID, params: &CodecParameters, options: Option<&Variant>) -> Result<Box<dyn Encoder<AudioEncoder>>> {
336        Ok(Box::new(OpusEncoder::new(codec_id, &params.try_into()?, options)?))
337    }
338}
339
340impl CodecBuilder<AudioEncoder> for OpusEncoderBuilder {
341    fn id(&self) -> CodecID {
342        CodecID::OPUS
343    }
344
345    fn name(&self) -> &'static str {
346        CODEC_NAME
347    }
348}
349
350impl CodecInformation for OpusEncoder {
351    fn id(&self) -> CodecID {
352        CodecID::OPUS
353    }
354
355    fn name(&self) -> &'static str {
356        CODEC_NAME
357    }
358}
359
360#[ctor]
361pub fn initialize() {
362    register_encoder(Arc::new(OpusEncoderBuilder), false);
363}