vpx_encode/
lib.rs

1//! Rust interface to libvpx encoder
2//!
3//! This crate provides a Rust API to use
4//! [libvpx](https://en.wikipedia.org/wiki/Libvpx) for encoding images.
5//!
6//! It it based entirely on code from [srs](https://crates.io/crates/srs).
7//! Compared to the original `srs`, this code has been simplified for use as a
8//! library and updated to add support for both the VP8 codec and (optionally)
9//! the VP9 codec.
10//!
11//! # Optional features
12//!
13//! Compile with the cargo feature `vp9` to enable support for the VP9 codec.
14//!
15//! # Example
16//!
17//! An example of using `vpx-encode` can be found in the [`record-screen`]()
18//! program. The source code for `record-screen` is in the [vpx-encode git
19//! repository]().
20//!
21//! # Contributing
22//!
23//! All contributions are appreciated.
24
25// vpx_sys is provided by the `env-libvpx-sys` crate
26
27#![cfg_attr(
28    feature = "backtrace",
29    feature(error_generic_member_access, provide_any)
30)]
31
32use std::{
33    mem::MaybeUninit,
34    os::raw::{c_int, c_uint, c_ulong},
35};
36
37#[cfg(feature = "backtrace")]
38use std::backtrace::Backtrace;
39use std::{ptr, slice};
40
41use thiserror::Error;
42
43#[cfg(feature = "vp9")]
44use vpx_sys::vp8e_enc_control_id::*;
45use vpx_sys::vpx_codec_cx_pkt_kind::VPX_CODEC_CX_FRAME_PKT;
46use vpx_sys::*;
47
48#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
49pub enum VideoCodecId {
50    VP8,
51    #[cfg(feature = "vp9")]
52    VP9,
53}
54
55impl Default for VideoCodecId {
56    #[cfg(not(feature = "vp9"))]
57    fn default() -> VideoCodecId {
58        VideoCodecId::VP8
59    }
60
61    #[cfg(feature = "vp9")]
62    fn default() -> VideoCodecId {
63        VideoCodecId::VP9
64    }
65}
66
67pub struct Encoder {
68    ctx: vpx_codec_ctx_t,
69    width: usize,
70    height: usize,
71}
72
73#[derive(Debug, Error)]
74#[error("VPX encode error: {msg}")]
75pub struct Error {
76    msg: String,
77    #[cfg(feature = "backtrace")]
78    #[backtrace]
79    backtrace: Backtrace,
80}
81
82impl From<String> for Error {
83    fn from(msg: String) -> Self {
84        Self {
85            msg,
86            #[cfg(feature = "backtrace")]
87            backtrace: Backtrace::capture(),
88        }
89    }
90}
91
92pub type Result<T> = std::result::Result<T, Error>;
93
94macro_rules! call_vpx {
95    ($x:expr) => {{
96        let result = unsafe { $x }; // original expression
97        let result_int = unsafe { std::mem::transmute::<_, i32>(result) };
98        // if result != VPX_CODEC_OK {
99        if result_int != 0 {
100            return Err(Error::from(format!(
101                "Function call failed (error code {}).",
102                result_int
103            )));
104        }
105        result
106    }};
107}
108
109macro_rules! call_vpx_ptr {
110    ($x:expr) => {{
111        let result = unsafe { $x }; // original expression
112        if result.is_null() {
113            return Err(Error::from("Bad pointer.".to_string()));
114        }
115        result
116    }};
117}
118
119impl Encoder {
120    pub fn new(config: Config) -> Result<Self> {
121        let i = match config.codec {
122            VideoCodecId::VP8 => call_vpx_ptr!(vpx_codec_vp8_cx()),
123            #[cfg(feature = "vp9")]
124            VideoCodecId::VP9 => call_vpx_ptr!(vpx_codec_vp9_cx()),
125        };
126
127        if config.width % 2 != 0 {
128            return Err(Error::from("Width must be divisible by 2".to_string()));
129        }
130        if config.height % 2 != 0 {
131            return Err(Error::from("Height must be divisible by 2".to_string()));
132        }
133
134        let c = MaybeUninit::zeroed();
135        let mut c = unsafe { c.assume_init() };
136        call_vpx!(vpx_codec_enc_config_default(i, &mut c, 0));
137
138        c.g_w = config.width;
139        c.g_h = config.height;
140        c.g_timebase.num = config.timebase[0];
141        c.g_timebase.den = config.timebase[1];
142        c.rc_target_bitrate = config.bitrate;
143
144        c.g_threads = 8;
145        c.g_error_resilient = VPX_ERROR_RESILIENT_DEFAULT;
146
147        let ctx = MaybeUninit::zeroed();
148        let mut ctx = unsafe { ctx.assume_init() };
149
150        match config.codec {
151            VideoCodecId::VP8 => {
152                call_vpx!(vpx_codec_enc_init_ver(
153                    &mut ctx,
154                    i,
155                    &c,
156                    0,
157                    vpx_sys::VPX_ENCODER_ABI_VERSION as i32
158                ));
159            }
160            #[cfg(feature = "vp9")]
161            VideoCodecId::VP9 => {
162                call_vpx!(vpx_codec_enc_init_ver(
163                    &mut ctx,
164                    i,
165                    &c,
166                    0,
167                    vpx_sys::VPX_ENCODER_ABI_VERSION as i32
168                ));
169                // set encoder internal speed settings
170                call_vpx!(vpx_codec_control_(
171                    &mut ctx,
172                    VP8E_SET_CPUUSED as _,
173                    6 as c_int
174                ));
175                // set row level multi-threading
176                call_vpx!(vpx_codec_control_(
177                    &mut ctx,
178                    VP9E_SET_ROW_MT as _,
179                    1 as c_int
180                ));
181            }
182        };
183
184        Ok(Self {
185            ctx,
186            width: config.width as usize,
187            height: config.height as usize,
188        })
189    }
190
191    pub fn encode(&mut self, pts: i64, data: &[u8]) -> Result<Packets> {
192        assert!(2 * data.len() >= 3 * self.width * self.height);
193
194        let image = MaybeUninit::zeroed();
195        let mut image = unsafe { image.assume_init() };
196
197        call_vpx_ptr!(vpx_img_wrap(
198            &mut image,
199            vpx_img_fmt::VPX_IMG_FMT_I420,
200            self.width as _,
201            self.height as _,
202            1,
203            data.as_ptr() as _,
204        ));
205
206        call_vpx!(vpx_codec_encode(
207            &mut self.ctx,
208            &image,
209            pts,
210            1, // Duration
211            0, // Flags
212            vpx_sys::VPX_DL_REALTIME as c_ulong,
213        ));
214
215        Ok(Packets {
216            ctx: &mut self.ctx,
217            iter: ptr::null(),
218        })
219    }
220
221    pub fn finish(mut self) -> Result<Finish> {
222        call_vpx!(vpx_codec_encode(
223            &mut self.ctx,
224            ptr::null(),
225            -1, // PTS
226            1,  // Duration
227            0,  // Flags
228            vpx_sys::VPX_DL_REALTIME as c_ulong,
229        ));
230
231        Ok(Finish {
232            enc: self,
233            iter: ptr::null(),
234        })
235    }
236}
237
238impl Drop for Encoder {
239    fn drop(&mut self) {
240        unsafe {
241            let result = vpx_codec_destroy(&mut self.ctx);
242            if result != vpx_sys::VPX_CODEC_OK {
243                eprintln!("failed to destroy vpx codec: {result:?}");
244            }
245        }
246    }
247}
248
249#[derive(Clone, Copy, Debug)]
250pub struct Frame<'a> {
251    /// Compressed data.
252    pub data: &'a [u8],
253    /// Whether the frame is a keyframe.
254    pub key: bool,
255    /// Presentation timestamp (in timebase units).
256    pub pts: i64,
257}
258
259#[derive(Clone, Copy, Debug)]
260pub struct Config {
261    /// The width (in pixels).
262    pub width: c_uint,
263    /// The height (in pixels).
264    pub height: c_uint,
265    /// The timebase numerator and denominator (in seconds).
266    pub timebase: [c_int; 2],
267    /// The target bitrate (in kilobits per second).
268    pub bitrate: c_uint,
269    /// The codec
270    pub codec: VideoCodecId,
271}
272
273pub struct Packets<'a> {
274    ctx: &'a mut vpx_codec_ctx_t,
275    iter: vpx_codec_iter_t,
276}
277
278impl<'a> Iterator for Packets<'a> {
279    type Item = Frame<'a>;
280    fn next(&mut self) -> Option<Self::Item> {
281        loop {
282            unsafe {
283                let pkt = vpx_codec_get_cx_data(self.ctx, &mut self.iter);
284                if pkt.is_null() {
285                    return None;
286                } else if (*pkt).kind == VPX_CODEC_CX_FRAME_PKT {
287                    let f = &(*pkt).data.frame;
288                    return Some(Frame {
289                        data: slice::from_raw_parts(f.buf as _, f.sz as usize),
290                        key: (f.flags & VPX_FRAME_IS_KEY) != 0,
291                        pts: f.pts,
292                    });
293                } else {
294                    // Ignore the packet.
295                }
296            }
297        }
298    }
299}
300
301pub struct Finish {
302    enc: Encoder,
303    iter: vpx_codec_iter_t,
304}
305
306impl Finish {
307    pub fn next(&mut self) -> Result<Option<Frame>> {
308        let mut tmp = Packets {
309            ctx: &mut self.enc.ctx,
310            iter: self.iter,
311        };
312
313        if let Some(packet) = tmp.next() {
314            self.iter = tmp.iter;
315            Ok(Some(packet))
316        } else {
317            call_vpx!(vpx_codec_encode(
318                tmp.ctx,
319                ptr::null(),
320                -1, // PTS
321                1,  // Duration
322                0,  // Flags
323                vpx_sys::VPX_DL_REALTIME as c_ulong,
324            ));
325
326            tmp.iter = ptr::null();
327            if let Some(packet) = tmp.next() {
328                self.iter = tmp.iter;
329                Ok(Some(packet))
330            } else {
331                Ok(None)
332            }
333        }
334    }
335}