videocall_cli/
video_encoder.rs

1use anyhow::{anyhow, Result};
2use std::mem::MaybeUninit;
3use std::os::raw::{c_int, c_ulong};
4use vpx_sys::*;
5
6macro_rules! vpx {
7    ($f:expr) => {{
8        let res = unsafe { $f };
9        let res_int = unsafe { std::mem::transmute::<vpx_sys::vpx_codec_err_t, i32>(res) };
10        if res_int != 0 {
11            return Err(anyhow!("vpx function error code ({}).", res_int));
12        }
13        res
14    }};
15}
16
17macro_rules! vpx_ptr {
18    ($f:expr) => {{
19        let res = unsafe { $f };
20        if res.is_null() {
21            return Err(anyhow!("vpx function returned null pointer."));
22        }
23        res
24    }};
25}
26
27pub struct VideoEncoderBuilder {
28    pub min_quantizer: u32,
29    pub max_quantizer: u32,
30    pub bitrate_kbps: u32,
31    pub fps: u32,
32    pub resolution: (u32, u32),
33    pub cpu_used: u32,
34    pub profile: u32,
35}
36
37impl VideoEncoderBuilder {
38    pub fn new(fps: u32, cpu_used: u8) -> Self {
39        Self {
40            bitrate_kbps: 500,
41            max_quantizer: 60,
42            min_quantizer: 40,
43            resolution: (640, 480),
44            fps,
45            cpu_used: cpu_used as u32,
46            profile: 0,
47        }
48    }
49}
50
51impl VideoEncoderBuilder {
52    pub fn set_resolution(mut self, width: u32, height: u32) -> Self {
53        self.resolution = (width, height);
54        self
55    }
56
57    pub fn build(&self) -> Result<VideoEncoder> {
58        let (width, height) = self.resolution;
59        if width % 2 != 0 || width == 0 {
60            return Err(anyhow!("Width must be divisible by 2"));
61        }
62        if height % 2 != 0 || height == 0 {
63            return Err(anyhow!("Height must be divisible by 2"));
64        }
65        let cfg_ptr = vpx_ptr!(vpx_codec_vp9_cx());
66        let mut cfg = unsafe { MaybeUninit::zeroed().assume_init() };
67        vpx!(vpx_codec_enc_config_default(cfg_ptr, &mut cfg, 0));
68
69        cfg.g_w = width;
70        cfg.g_h = height;
71        cfg.g_timebase.num = 1;
72        cfg.g_timebase.den = self.fps as c_int;
73        cfg.rc_target_bitrate = self.bitrate_kbps;
74        cfg.rc_min_quantizer = self.min_quantizer;
75        cfg.rc_max_quantizer = self.max_quantizer;
76        cfg.g_threads = 2;
77        cfg.g_lag_in_frames = 1;
78        cfg.g_error_resilient = VPX_ERROR_RESILIENT_DEFAULT;
79        cfg.g_pass = vpx_enc_pass::VPX_RC_ONE_PASS;
80        cfg.g_profile = self.profile;
81        cfg.rc_end_usage = vpx_rc_mode::VPX_VBR;
82        cfg.kf_max_dist = 150;
83        cfg.kf_min_dist = 150;
84        cfg.kf_mode = vpx_kf_mode::VPX_KF_AUTO;
85
86        let ctx = MaybeUninit::zeroed();
87        let mut ctx = unsafe { ctx.assume_init() };
88
89        vpx!(vpx_codec_enc_init_ver(
90            &mut ctx,
91            cfg_ptr,
92            &cfg,
93            0,
94            VPX_ENCODER_ABI_VERSION as i32
95        ));
96        unsafe {
97            vpx_codec_control_(&mut ctx, vp8e_enc_control_id::VP8E_SET_CPUUSED as c_int, 5);
98            vpx_codec_control_(
99                &mut ctx,
100                vp8e_enc_control_id::VP9E_SET_TILE_COLUMNS as c_int,
101                4,
102            );
103            vpx_codec_control_(&mut ctx, vp8e_enc_control_id::VP9E_SET_ROW_MT as c_int, 1);
104            vpx_codec_control_(
105                &mut ctx,
106                vp8e_enc_control_id::VP9E_SET_FRAME_PARALLEL_DECODING as c_int,
107                1,
108            );
109            // vpx_codec_control_(&mut ctx, vp8e_enc_control_id::VP9E_SET_AQ_MODE as c_int, 3);
110        }
111        Ok(VideoEncoder {
112            ctx,
113            cfg,
114            width: self.resolution.0,
115            height: self.resolution.1,
116        })
117    }
118}
119
120pub struct VideoEncoder {
121    ctx: vpx_codec_ctx_t,
122    cfg: vpx_codec_enc_cfg_t,
123    width: u32,
124    height: u32,
125}
126
127impl VideoEncoder {
128    pub fn update_bitrate_kbps(&mut self, bitrate: u32) -> anyhow::Result<()> {
129        self.cfg.rc_target_bitrate = bitrate;
130        vpx!(vpx_codec_enc_config_set(&mut self.ctx, &self.cfg));
131        Ok(())
132    }
133
134    pub fn encode(&mut self, pts: i64, data: &[u8]) -> anyhow::Result<Frames> {
135        let image = MaybeUninit::zeroed();
136        let mut image = unsafe { image.assume_init() };
137
138        vpx_ptr!(vpx_img_wrap(
139            &mut image,
140            vpx_img_fmt::VPX_IMG_FMT_I420,
141            self.width as _,
142            self.height as _,
143            1,
144            data.as_ptr() as _,
145        ));
146
147        let flags: i64 = 0;
148
149        vpx!(vpx_codec_encode(
150            &mut self.ctx,
151            &image,
152            pts,
153            1,     // Duration
154            flags, // Flags
155            VPX_DL_REALTIME as c_ulong,
156        ));
157
158        Ok(Frames {
159            ctx: &mut self.ctx,
160            iter: std::ptr::null(),
161        })
162    }
163}
164
165#[derive(Clone, Copy, Debug)]
166pub struct Frame<'a> {
167    /// Compressed data.
168    pub data: &'a [u8],
169    /// Whether the frame is a keyframe.
170    pub key: bool,
171    /// Presentation timestamp (in timebase units).
172    pub pts: i64,
173}
174
175pub struct Frames<'a> {
176    ctx: &'a mut vpx_codec_ctx_t,
177    iter: vpx_codec_iter_t,
178}
179
180impl<'a> Iterator for Frames<'a> {
181    type Item = Frame<'a>;
182    #[allow(clippy::unnecessary_cast)]
183    fn next(&mut self) -> Option<Self::Item> {
184        loop {
185            unsafe {
186                let pkt = vpx_codec_get_cx_data(self.ctx, &mut self.iter);
187                if pkt.is_null() {
188                    return None;
189                } else if (*pkt).kind == vpx_codec_cx_pkt_kind::VPX_CODEC_CX_FRAME_PKT {
190                    let f = &(*pkt).data.frame;
191                    return Some(Frame {
192                        data: std::slice::from_raw_parts(f.buf as _, f.sz as usize),
193                        key: (f.flags & VPX_FRAME_IS_KEY) != 0,
194                        pts: f.pts,
195                    });
196                }
197            }
198        }
199    }
200}