media_codec/
encoder.rs

1use std::{
2    any::TypeId,
3    collections::HashMap,
4    mem,
5    sync::{Arc, LazyLock, RwLock},
6};
7
8#[cfg(feature = "audio")]
9use media_core::audio::AudioFrameDescriptor;
10#[cfg(feature = "video")]
11use media_core::video::VideoFrameDescriptor;
12use media_core::{
13    buffer::BufferPool,
14    error::Error,
15    frame::{Frame, SharedFrame},
16    invalid_param_error,
17    rational::Rational64,
18    unsupported_error,
19    variant::Variant,
20    MediaType, Result,
21};
22
23#[cfg(feature = "audio")]
24use crate::AudioParameters;
25#[cfg(feature = "video")]
26use crate::VideoParameters;
27use crate::{
28    find_codec, find_codec_by_name, packet::Packet, register_codec, Codec, CodecBuilder, CodecID, CodecList, CodecParameters, CodecParametersType,
29    CodecSpec, CodecType, LazyCodecList, MediaParametersType,
30};
31
32#[derive(Clone, Debug, Default)]
33pub struct EncoderParameters {
34    pub bit_rate: Option<u64>,
35    pub profile: Option<i32>,
36    pub level: Option<i32>,
37    pub use_pool: Option<bool>,
38}
39
40impl EncoderParameters {
41    fn update(&mut self, other: &EncoderParameters) {
42        if other.bit_rate.is_some() {
43            self.bit_rate = other.bit_rate;
44        }
45        if other.profile.is_some() {
46            self.profile = other.profile;
47        }
48        if other.level.is_some() {
49            self.level = other.level;
50        }
51    }
52
53    fn update_with_option(&mut self, key: &str, value: &Variant) {
54        match key {
55            "bit_rate" => self.bit_rate = value.get_uint64(),
56            "profile" => self.profile = value.get_int32(),
57            "level" => self.level = value.get_int32(),
58            _ => {}
59        }
60    }
61}
62
63impl TryFrom<&CodecParametersType> for EncoderParameters {
64    type Error = Error;
65
66    fn try_from(params: &CodecParametersType) -> Result<Self> {
67        match params {
68            CodecParametersType::Encoder(params) => Ok(params.clone()),
69            _ => Err(invalid_param_error!(params)),
70        }
71    }
72}
73
74#[cfg(feature = "audio")]
75#[derive(Clone, Debug, Default)]
76pub struct AudioEncoderParameters {
77    pub audio: AudioParameters,
78    pub encoder: EncoderParameters,
79}
80
81#[cfg(feature = "audio")]
82#[allow(unreachable_patterns)]
83impl TryFrom<&CodecParameters> for AudioEncoderParameters {
84    type Error = Error;
85
86    fn try_from(params: &CodecParameters) -> Result<Self> {
87        Ok(Self {
88            audio: match &params.media {
89                MediaParametersType::Audio(params) => params.clone(),
90                _ => return Err(invalid_param_error!(params)),
91            },
92            encoder: match &params.codec {
93                CodecParametersType::Encoder(params) => params.clone(),
94                _ => return Err(invalid_param_error!(params)),
95            },
96        })
97    }
98}
99
100#[cfg(feature = "audio")]
101#[derive(Clone, Debug)]
102pub struct AudioEncoder {
103    pub audio: AudioParameters,
104    pub encoder: EncoderParameters,
105    // audio encoder specific configuration
106    pub frame_size: Option<u32>,
107    pub delay: Option<u32>,
108}
109
110#[cfg(feature = "audio")]
111impl CodecSpec for AudioEncoder {
112    type FrameDescriptor = AudioFrameDescriptor;
113
114    fn media_type() -> MediaType {
115        MediaType::Audio
116    }
117
118    fn codec_type() -> CodecType {
119        CodecType::Encoder
120    }
121
122    fn from_parameters(params: &CodecParameters) -> Result<Self> {
123        Ok(Self {
124            audio: (&params.media).try_into()?,
125            encoder: (&params.codec).try_into()?,
126            frame_size: None,
127            delay: None,
128        })
129    }
130
131    fn configure(&mut self, params: &CodecParameters) -> Result<()> {
132        let audio_params = (&params.media).try_into()?;
133        let encoder_params = (&params.codec).try_into()?;
134        self.audio.update(&audio_params);
135        self.encoder.update(&encoder_params);
136        Ok(())
137    }
138
139    fn configure_with_option(&mut self, key: &str, value: &Variant) -> Result<()> {
140        self.audio.update_with_option(key, value);
141        self.encoder.update_with_option(key, value);
142
143        match key {
144            "frame_size" => self.frame_size = value.get_uint32(),
145            "delay" => self.delay = value.get_uint32(),
146            _ => {}
147        }
148
149        Ok(())
150    }
151}
152
153#[cfg(feature = "video")]
154#[derive(Clone, Debug, Default)]
155pub struct VideoEncoderParameters {
156    pub video: VideoParameters,
157    pub encoder: EncoderParameters,
158}
159
160#[cfg(feature = "video")]
161#[allow(unreachable_patterns)]
162impl TryFrom<&CodecParameters> for VideoEncoderParameters {
163    type Error = Error;
164
165    #[allow(unreachable_patterns)]
166    fn try_from(params: &CodecParameters) -> Result<Self> {
167        Ok(Self {
168            video: match &params.media {
169                MediaParametersType::Video(params) => params.clone(),
170                _ => return Err(invalid_param_error!(params)),
171            },
172            encoder: match &params.codec {
173                CodecParametersType::Encoder(params) => params.clone(),
174                _ => return Err(invalid_param_error!(params)),
175            },
176        })
177    }
178}
179
180#[cfg(feature = "video")]
181#[derive(Clone, Debug)]
182pub struct VideoEncoder {
183    pub video: VideoParameters,
184    pub encoder: EncoderParameters,
185}
186
187#[cfg(feature = "video")]
188impl CodecSpec for VideoEncoder {
189    type FrameDescriptor = VideoFrameDescriptor;
190
191    fn media_type() -> MediaType {
192        MediaType::Video
193    }
194
195    fn codec_type() -> CodecType {
196        CodecType::Encoder
197    }
198
199    #[allow(unreachable_patterns)]
200    fn from_parameters(params: &CodecParameters) -> Result<Self> {
201        Ok(Self {
202            video: match &params.media {
203                MediaParametersType::Video(params) => params.clone(),
204                _ => return Err(invalid_param_error!(params)),
205            },
206            encoder: match &params.codec {
207                CodecParametersType::Encoder(params) => params.clone(),
208                _ => return Err(invalid_param_error!(params)),
209            },
210        })
211    }
212
213    #[allow(unreachable_patterns)]
214    fn configure(&mut self, params: &CodecParameters) -> Result<()> {
215        let video_params = (&params.media).try_into()?;
216        let encoder_params = (&params.codec).try_into()?;
217        self.video.update(&video_params);
218        self.encoder.update(&encoder_params);
219        Ok(())
220    }
221
222    fn configure_with_option(&mut self, key: &str, value: &Variant) -> Result<()> {
223        self.video.update_with_option(key, value);
224        self.encoder.update_with_option(key, value);
225        Ok(())
226    }
227}
228
229pub trait Encoder<T: CodecSpec>: Codec<T> + Send + Sync {
230    fn init(&mut self, _config: &T) -> Result<()> {
231        Ok(())
232    }
233    fn send_frame(&mut self, config: &T, pool: Option<&Arc<BufferPool>>, frame: SharedFrame<Frame<'static, T::FrameDescriptor>>) -> Result<()>;
234    fn receive_packet(&mut self, config: &T, pool: Option<&Arc<BufferPool>>) -> Result<Packet<'static>>;
235    fn flush(&mut self, config: &T) -> Result<()>;
236}
237
238pub trait EncoderBuilder<T: CodecSpec>: CodecBuilder<T> {
239    fn new_encoder(&self, id: CodecID, params: &CodecParameters, options: Option<&Variant>) -> Result<Box<dyn Encoder<T>>>;
240}
241
242pub struct EncoderContext<T: CodecSpec> {
243    pub config: T,
244    pub time_base: Option<Rational64>,
245    encoder: Box<dyn Encoder<T>>,
246    pool: Option<Arc<BufferPool>>,
247}
248
249#[cfg(feature = "audio")]
250static AUDIO_ENCODER_LIST: LazyCodecList<AudioEncoder> = LazyLock::new(|| {
251    RwLock::new(CodecList::<AudioEncoder> {
252        codecs: HashMap::new(),
253    })
254});
255
256#[cfg(feature = "video")]
257static VIDEO_ENCODER_LIST: LazyCodecList<VideoEncoder> = LazyLock::new(|| {
258    RwLock::new(CodecList::<VideoEncoder> {
259        codecs: HashMap::new(),
260    })
261});
262
263pub fn register_encoder<T: CodecSpec>(builder: Arc<dyn EncoderBuilder<T>>, default: bool) -> Result<()> {
264    match TypeId::of::<T>() {
265        #[cfg(feature = "audio")]
266        id if id == TypeId::of::<AudioEncoder>() => {
267            let builder = unsafe { mem::transmute::<Arc<dyn EncoderBuilder<T>>, Arc<dyn CodecBuilder<AudioEncoder>>>(builder) };
268            register_codec(&AUDIO_ENCODER_LIST, builder, default)
269        }
270        #[cfg(feature = "video")]
271        id if id == TypeId::of::<VideoEncoder>() => {
272            let builder = unsafe { mem::transmute::<Arc<dyn EncoderBuilder<T>>, Arc<dyn CodecBuilder<VideoEncoder>>>(builder) };
273            register_codec(&VIDEO_ENCODER_LIST, builder, default)
274        }
275        _ => Err(unsupported_error!("codec parameters type")),
276    }
277}
278
279pub(crate) fn find_encoder<T: CodecSpec>(id: CodecID) -> Result<Arc<dyn EncoderBuilder<T>>> {
280    match TypeId::of::<T>() {
281        #[cfg(feature = "audio")]
282        type_id if type_id == TypeId::of::<AudioEncoder>() => {
283            let builder = find_codec(&AUDIO_ENCODER_LIST, id)?;
284            unsafe { Ok(mem::transmute::<Arc<dyn CodecBuilder<AudioEncoder>>, Arc<dyn EncoderBuilder<T>>>(builder)) }
285        }
286        #[cfg(feature = "video")]
287        type_id if type_id == TypeId::of::<VideoEncoder>() => {
288            let builder = find_codec(&VIDEO_ENCODER_LIST, id)?;
289            unsafe { Ok(mem::transmute::<Arc<dyn CodecBuilder<VideoEncoder>>, Arc<dyn EncoderBuilder<T>>>(builder)) }
290        }
291        _ => Err(unsupported_error!("codec parameters type")),
292    }
293}
294
295pub(crate) fn find_encoder_by_name<T: CodecSpec>(name: &str) -> Result<Arc<dyn EncoderBuilder<T>>> {
296    match TypeId::of::<T>() {
297        #[cfg(feature = "audio")]
298        id if id == TypeId::of::<AudioEncoder>() => {
299            let builder = find_codec_by_name(&AUDIO_ENCODER_LIST, name)?;
300            unsafe { Ok(mem::transmute::<Arc<dyn CodecBuilder<AudioEncoder>>, Arc<dyn EncoderBuilder<T>>>(builder)) }
301        }
302        #[cfg(feature = "video")]
303        id if id == TypeId::of::<VideoEncoder>() => {
304            let builder = find_codec_by_name(&VIDEO_ENCODER_LIST, name)?;
305            unsafe { Ok(mem::transmute::<Arc<dyn CodecBuilder<VideoEncoder>>, Arc<dyn EncoderBuilder<T>>>(builder)) }
306        }
307        _ => Err(unsupported_error!("codec parameters type")),
308    }
309}
310
311impl<T: CodecSpec> EncoderContext<T> {
312    fn new_with_builder(builder: Arc<dyn EncoderBuilder<T>>, params: &CodecParameters, options: Option<&Variant>) -> Result<Self> {
313        let encoder = builder.new_encoder(builder.id(), params, options)?;
314        let config = T::from_parameters(params)?;
315
316        let buffer_pool = match &params.codec {
317            CodecParametersType::Encoder(encoder_params) => {
318                if encoder_params.use_pool.unwrap_or(false) {
319                    Some(BufferPool::new(0))
320                } else {
321                    None
322                }
323            }
324            _ => return Err(invalid_param_error!(params)),
325        };
326
327        Ok(Self {
328            config,
329            time_base: None,
330            encoder,
331            pool: buffer_pool,
332        })
333    }
334
335    pub fn from_codec_id(id: CodecID, params: &CodecParameters, options: Option<&Variant>) -> Result<Self> {
336        Self::new_with_builder(find_encoder(id)?, params, options)
337    }
338
339    pub fn from_codec_name(name: &str, params: &CodecParameters, options: Option<&Variant>) -> Result<Self> {
340        Self::new_with_builder(find_encoder_by_name(name)?, params, options)
341    }
342
343    pub fn configure(&mut self, params: Option<&CodecParameters>, options: Option<&Variant>) -> Result<()> {
344        if let Some(params) = params {
345            self.config.configure(params)?;
346        }
347        self.encoder.configure(params, options)
348    }
349
350    pub fn set_option(&mut self, key: &str, value: &Variant) -> Result<()> {
351        self.encoder.set_option(key, value)
352    }
353
354    pub fn send_frame(&mut self, frame: SharedFrame<Frame<'static, T::FrameDescriptor>>) -> Result<()> {
355        self.encoder.send_frame(&self.config, self.pool.as_ref(), frame)
356    }
357
358    pub fn receive_packet(&mut self) -> Result<Packet<'static>> {
359        let mut packet = self.encoder.receive_packet(&self.config, self.pool.as_ref())?;
360
361        packet.time_base = packet.time_base.or(self.time_base);
362
363        Ok(packet)
364    }
365}