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