612 lines
18 KiB
Rust
612 lines
18 KiB
Rust
// https://github.com/astraw/vpx-encode
|
||
// https://github.com/astraw/env-libvpx-sys
|
||
// https://github.com/rust-av/vpx-rs/blob/master/src/decoder.rs
|
||
|
||
use hbb_common::anyhow::{anyhow, Context};
|
||
use hbb_common::message_proto::{test_delay, Message, VP9s, VideoFrame, VP9};
|
||
use hbb_common::ResultType;
|
||
|
||
use crate::codec::EncoderApi;
|
||
use crate::STRIDE_ALIGN;
|
||
|
||
use super::vpx::{vp8e_enc_control_id::*, vpx_codec_err_t::*, *};
|
||
use std::os::raw::{c_int, c_uint};
|
||
use std::{ptr, slice};
|
||
|
||
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
|
||
pub enum VpxVideoCodecId {
|
||
VP8,
|
||
VP9,
|
||
}
|
||
|
||
impl Default for VpxVideoCodecId {
|
||
fn default() -> VpxVideoCodecId {
|
||
VpxVideoCodecId::VP9
|
||
}
|
||
}
|
||
|
||
pub struct VpxEncoder {
|
||
ctx: vpx_codec_ctx_t,
|
||
format: VpxVideoCodecId,
|
||
width: usize,
|
||
height: usize,
|
||
}
|
||
|
||
pub struct VpxDecoder {
|
||
ctx: vpx_codec_ctx_t,
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
pub enum Error {
|
||
FailedCall(String),
|
||
BadPtr(String),
|
||
}
|
||
|
||
impl std::fmt::Display for Error {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::result::Result<(), std::fmt::Error> {
|
||
write!(f, "{:?}", self)
|
||
}
|
||
}
|
||
|
||
impl std::error::Error for Error {}
|
||
|
||
pub type Result<T> = std::result::Result<T, Error>;
|
||
|
||
macro_rules! call_vpx {
|
||
($x:expr) => {{
|
||
let result = unsafe { $x }; // original expression
|
||
let result_int = unsafe { std::mem::transmute::<_, i32>(result) };
|
||
if result_int != 0 {
|
||
return Err(Error::FailedCall(format!(
|
||
"errcode={} {}:{}:{}:{}",
|
||
result_int,
|
||
module_path!(),
|
||
file!(),
|
||
line!(),
|
||
column!()
|
||
))
|
||
.into());
|
||
}
|
||
result
|
||
}};
|
||
}
|
||
|
||
macro_rules! call_vpx_ptr {
|
||
($x:expr) => {{
|
||
let result = unsafe { $x }; // original expression
|
||
let result_int = unsafe { std::mem::transmute::<_, isize>(result) };
|
||
if result_int == 0 {
|
||
return Err(Error::BadPtr(format!(
|
||
"errcode={} {}:{}:{}:{}",
|
||
result_int,
|
||
module_path!(),
|
||
file!(),
|
||
line!(),
|
||
column!()
|
||
))
|
||
.into());
|
||
}
|
||
result
|
||
}};
|
||
}
|
||
|
||
impl EncoderApi for VpxEncoder {
|
||
fn new(cfg: crate::codec::EncoderCfg) -> ResultType<Self>
|
||
where
|
||
Self: Sized,
|
||
{
|
||
match cfg {
|
||
crate::codec::EncoderCfg::VPX(config) => {
|
||
let format;
|
||
let i;
|
||
if cfg!(feature = "VP8") {
|
||
i = match config.codec {
|
||
VpxVideoCodecId::VP8 => call_vpx_ptr!(vpx_codec_vp8_cx()),
|
||
VpxVideoCodecId::VP9 => call_vpx_ptr!(vpx_codec_vp9_cx()),
|
||
};
|
||
format = config.codec;
|
||
} else {
|
||
i = call_vpx_ptr!(vpx_codec_vp9_cx());
|
||
format = VpxVideoCodecId::VP9;
|
||
}
|
||
let mut c = unsafe { std::mem::MaybeUninit::zeroed().assume_init() };
|
||
call_vpx!(vpx_codec_enc_config_default(i, &mut c, 0));
|
||
|
||
// https://www.webmproject.org/docs/encoder-parameters/
|
||
// default: c.rc_min_quantizer = 0, c.rc_max_quantizer = 63
|
||
// try rc_resize_allowed later
|
||
|
||
c.g_w = config.width;
|
||
c.g_h = config.height;
|
||
c.g_timebase.num = config.timebase[0];
|
||
c.g_timebase.den = config.timebase[1];
|
||
c.rc_target_bitrate = config.bitrate;
|
||
c.rc_undershoot_pct = 95;
|
||
c.rc_dropframe_thresh = 25;
|
||
c.g_threads = if config.num_threads == 0 {
|
||
num_cpus::get() as _
|
||
} else {
|
||
config.num_threads
|
||
};
|
||
c.g_error_resilient = VPX_ERROR_RESILIENT_DEFAULT;
|
||
// https://developers.google.com/media/vp9/bitrate-modes/
|
||
// Constant Bitrate mode (CBR) is recommended for live streaming with VP9.
|
||
c.rc_end_usage = vpx_rc_mode::VPX_CBR;
|
||
// c.kf_min_dist = 0;
|
||
// c.kf_max_dist = 999999;
|
||
c.kf_mode = vpx_kf_mode::VPX_KF_DISABLED; // reduce bandwidth a lot
|
||
|
||
/*
|
||
VPX encoder支持two-pass encode，这是为了rate control的。
|
||
对于两遍编码，就是需要整个编码过程做两次，第一次会得到一些新的控制参数来进行第二遍的编码，
|
||
这样可以在相同的bitrate下得到最好的PSNR
|
||
*/
|
||
|
||
let mut ctx = Default::default();
|
||
call_vpx!(vpx_codec_enc_init_ver(
|
||
&mut ctx,
|
||
i,
|
||
&c,
|
||
0,
|
||
VPX_ENCODER_ABI_VERSION as _
|
||
));
|
||
|
||
if config.codec == VpxVideoCodecId::VP9 {
|
||
// set encoder internal speed settings
|
||
// in ffmpeg, it is --speed option
|
||
/*
|
||
set to 0 or a positive value 1-16, the codec will try to adapt its
|
||
complexity depending on the time it spends encoding. Increasing this
|
||
number will make the speed go up and the quality go down.
|
||
Negative values mean strict enforcement of this
|
||
while positive values are adaptive
|
||
*/
|
||
/* https://developers.google.com/media/vp9/live-encoding
|
||
Speed 5 to 8 should be used for live / real-time encoding.
|
||
Lower numbers (5 or 6) are higher quality but require more CPU power.
|
||
Higher numbers (7 or 8) will be lower quality but more manageable for lower latency
|
||
use cases and also for lower CPU power devices such as mobile.
|
||
*/
|
||
call_vpx!(vpx_codec_control_(&mut ctx, VP8E_SET_CPUUSED as _, 7,));
|
||
// set row level multi-threading
|
||
/*
|
||
as some people in comments and below have already commented,
|
||
more recent versions of libvpx support -row-mt 1 to enable tile row
|
||
multi-threading. This can increase the number of tiles by up to 4x in VP9
|
||
(since the max number of tile rows is 4, regardless of video height).
|
||
To enable this, use -tile-rows N where N is the number of tile rows in
|
||
log2 units (so -tile-rows 1 means 2 tile rows and -tile-rows 2 means 4 tile
|
||
rows). The total number of active threads will then be equal to
|
||
$tile_rows * $tile_columns
|
||
*/
|
||
call_vpx!(vpx_codec_control_(
|
||
&mut ctx,
|
||
VP9E_SET_ROW_MT as _,
|
||
1 as c_int
|
||
));
|
||
|
||
call_vpx!(vpx_codec_control_(
|
||
&mut ctx,
|
||
VP9E_SET_TILE_COLUMNS as _,
|
||
4 as c_int
|
||
));
|
||
}
|
||
|
||
Ok(Self {
|
||
ctx,
|
||
format,
|
||
width: config.width as _,
|
||
height: config.height as _,
|
||
})
|
||
}
|
||
_ => Err(anyhow!("encoder type mismatch")),
|
||
}
|
||
}
|
||
|
||
fn encode_to_message(&mut self, frame: &[u8], ms: i64) -> ResultType<Message> {
|
||
let mut frames = Vec::new();
|
||
for ref frame in self
|
||
.encode(ms, frame, STRIDE_ALIGN)
|
||
.with_context(|| "Failed to encode")?
|
||
{
|
||
frames.push(VpxEncoder::create_frame(frame));
|
||
}
|
||
for ref frame in self.flush().with_context(|| "Failed to flush")? {
|
||
frames.push(VpxEncoder::create_frame(frame));
|
||
}
|
||
|
||
// to-do: flush periodically, e.g. 1 second
|
||
if frames.len() > 0 {
|
||
Ok(VpxEncoder::create_msg(frames))
|
||
} else {
|
||
Err(anyhow!("no valid frame"))
|
||
}
|
||
}
|
||
|
||
fn use_yuv(&self) -> bool {
|
||
true
|
||
}
|
||
|
||
fn set_bitrate(&mut self, bitrate: u32) -> ResultType<()> {
|
||
let mut new_enc_cfg = unsafe { *self.ctx.config.enc.to_owned() };
|
||
new_enc_cfg.rc_target_bitrate = bitrate;
|
||
call_vpx!(vpx_codec_enc_config_set(&mut self.ctx, &new_enc_cfg));
|
||
return Ok(());
|
||
}
|
||
|
||
fn get_codec_format(&self) -> test_delay::CodecFormat {
|
||
match self.format {
|
||
VpxVideoCodecId::VP8 => test_delay::CodecFormat::VP8,
|
||
VpxVideoCodecId::VP9 => test_delay::CodecFormat::VP9,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl VpxEncoder {
|
||
pub fn encode(&mut self, pts: i64, data: &[u8], stride_align: usize) -> Result<EncodeFrames> {
|
||
assert!(2 * data.len() >= 3 * self.width * self.height);
|
||
|
||
let mut image = Default::default();
|
||
call_vpx_ptr!(vpx_img_wrap(
|
||
&mut image,
|
||
vpx_img_fmt::VPX_IMG_FMT_I420,
|
||
self.width as _,
|
||
self.height as _,
|
||
stride_align as _,
|
||
data.as_ptr() as _,
|
||
));
|
||
|
||
call_vpx!(vpx_codec_encode(
|
||
&mut self.ctx,
|
||
&image,
|
||
pts as _,
|
||
1, // Duration
|
||
0, // Flags
|
||
VPX_DL_REALTIME as _,
|
||
));
|
||
|
||
Ok(EncodeFrames {
|
||
ctx: &mut self.ctx,
|
||
iter: ptr::null(),
|
||
})
|
||
}
|
||
|
||
/// Notify the encoder to return any pending packets
|
||
pub fn flush(&mut self) -> Result<EncodeFrames> {
|
||
call_vpx!(vpx_codec_encode(
|
||
&mut self.ctx,
|
||
ptr::null(),
|
||
-1, // PTS
|
||
1, // Duration
|
||
0, // Flags
|
||
VPX_DL_REALTIME as _,
|
||
));
|
||
|
||
Ok(EncodeFrames {
|
||
ctx: &mut self.ctx,
|
||
iter: ptr::null(),
|
||
})
|
||
}
|
||
|
||
#[inline]
|
||
fn create_msg(vp9s: Vec<VP9>) -> Message {
|
||
let mut msg_out = Message::new();
|
||
let mut vf = VideoFrame::new();
|
||
vf.set_vp9s(VP9s {
|
||
frames: vp9s.into(),
|
||
..Default::default()
|
||
});
|
||
msg_out.set_video_frame(vf);
|
||
msg_out
|
||
}
|
||
|
||
#[inline]
|
||
fn create_frame(frame: &EncodeFrame) -> VP9 {
|
||
VP9 {
|
||
data: frame.data.to_vec(),
|
||
key: frame.key,
|
||
pts: frame.pts,
|
||
..Default::default()
|
||
}
|
||
}
|
||
}
|
||
|
||
impl Drop for VpxEncoder {
|
||
fn drop(&mut self) {
|
||
unsafe {
|
||
let result = vpx_codec_destroy(&mut self.ctx);
|
||
if result != VPX_CODEC_OK {
|
||
panic!("failed to destroy vpx codec");
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Copy, Debug)]
|
||
pub struct EncodeFrame<'a> {
|
||
/// Compressed data.
|
||
pub data: &'a [u8],
|
||
/// Whether the frame is a keyframe.
|
||
pub key: bool,
|
||
/// Presentation timestamp (in timebase units).
|
||
pub pts: i64,
|
||
}
|
||
|
||
#[derive(Clone, Copy, Debug)]
|
||
pub struct VpxEncoderConfig {
|
||
/// The width (in pixels).
|
||
pub width: c_uint,
|
||
/// The height (in pixels).
|
||
pub height: c_uint,
|
||
/// The timebase numerator and denominator (in seconds).
|
||
pub timebase: [c_int; 2],
|
||
/// The target bitrate (in kilobits per second).
|
||
pub bitrate: c_uint,
|
||
/// The codec
|
||
pub codec: VpxVideoCodecId,
|
||
pub num_threads: u32,
|
||
}
|
||
|
||
#[derive(Clone, Copy, Debug)]
|
||
pub struct VpxDecoderConfig {
|
||
pub codec: VpxVideoCodecId,
|
||
pub num_threads: u32,
|
||
}
|
||
|
||
pub struct EncodeFrames<'a> {
|
||
ctx: &'a mut vpx_codec_ctx_t,
|
||
iter: vpx_codec_iter_t,
|
||
}
|
||
|
||
impl<'a> Iterator for EncodeFrames<'a> {
|
||
type Item = EncodeFrame<'a>;
|
||
fn next(&mut self) -> Option<Self::Item> {
|
||
loop {
|
||
unsafe {
|
||
let pkt = vpx_codec_get_cx_data(self.ctx, &mut self.iter);
|
||
if pkt.is_null() {
|
||
return None;
|
||
} else if (*pkt).kind == vpx_codec_cx_pkt_kind::VPX_CODEC_CX_FRAME_PKT {
|
||
let f = &(*pkt).data.frame;
|
||
return Some(Self::Item {
|
||
data: slice::from_raw_parts(f.buf as _, f.sz as _),
|
||
key: (f.flags & VPX_FRAME_IS_KEY) != 0,
|
||
pts: f.pts,
|
||
});
|
||
} else {
|
||
// Ignore the packet.
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
impl VpxDecoder {
|
||
/// Create a new decoder
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// The function may fail if the underlying libvpx does not provide
|
||
/// the VP9 decoder.
|
||
pub fn new(config: VpxDecoderConfig) -> Result<Self> {
|
||
// This is sound because `vpx_codec_ctx` is a repr(C) struct without any field that can
|
||
// cause UB if uninitialized.
|
||
let i;
|
||
if cfg!(feature = "VP8") {
|
||
i = match config.codec {
|
||
VpxVideoCodecId::VP8 => call_vpx_ptr!(vpx_codec_vp8_dx()),
|
||
VpxVideoCodecId::VP9 => call_vpx_ptr!(vpx_codec_vp9_dx()),
|
||
};
|
||
} else {
|
||
i = call_vpx_ptr!(vpx_codec_vp9_dx());
|
||
}
|
||
let mut ctx = Default::default();
|
||
let cfg = vpx_codec_dec_cfg_t {
|
||
threads: if config.num_threads == 0 {
|
||
num_cpus::get() as _
|
||
} else {
|
||
config.num_threads
|
||
},
|
||
w: 0,
|
||
h: 0,
|
||
};
|
||
/*
|
||
unsafe {
|
||
println!("{}", vpx_codec_get_caps(i));
|
||
}
|
||
*/
|
||
call_vpx!(vpx_codec_dec_init_ver(
|
||
&mut ctx,
|
||
i,
|
||
&cfg,
|
||
0,
|
||
VPX_DECODER_ABI_VERSION as _,
|
||
));
|
||
Ok(Self { ctx })
|
||
}
|
||
|
||
pub fn decode2rgb(&mut self, data: &[u8], rgba: bool) -> Result<Vec<u8>> {
|
||
let mut img = Image::new();
|
||
for frame in self.decode(data)? {
|
||
drop(img);
|
||
img = frame;
|
||
}
|
||
for frame in self.flush()? {
|
||
drop(img);
|
||
img = frame;
|
||
}
|
||
if img.is_null() {
|
||
Ok(Vec::new())
|
||
} else {
|
||
let mut out = Default::default();
|
||
img.rgb(1, rgba, &mut out);
|
||
Ok(out)
|
||
}
|
||
}
|
||
|
||
/// Feed some compressed data to the encoder
|
||
///
|
||
/// The `data` slice is sent to the decoder
|
||
///
|
||
/// It matches a call to `vpx_codec_decode`.
|
||
pub fn decode(&mut self, data: &[u8]) -> Result<DecodeFrames> {
|
||
call_vpx!(vpx_codec_decode(
|
||
&mut self.ctx,
|
||
data.as_ptr(),
|
||
data.len() as _,
|
||
ptr::null_mut(),
|
||
0,
|
||
));
|
||
|
||
Ok(DecodeFrames {
|
||
ctx: &mut self.ctx,
|
||
iter: ptr::null(),
|
||
})
|
||
}
|
||
|
||
/// Notify the decoder to return any pending frame
|
||
pub fn flush(&mut self) -> Result<DecodeFrames> {
|
||
call_vpx!(vpx_codec_decode(
|
||
&mut self.ctx,
|
||
ptr::null(),
|
||
0,
|
||
ptr::null_mut(),
|
||
0
|
||
));
|
||
Ok(DecodeFrames {
|
||
ctx: &mut self.ctx,
|
||
iter: ptr::null(),
|
||
})
|
||
}
|
||
}
|
||
|
||
impl Drop for VpxDecoder {
|
||
fn drop(&mut self) {
|
||
unsafe {
|
||
let result = vpx_codec_destroy(&mut self.ctx);
|
||
if result != VPX_CODEC_OK {
|
||
panic!("failed to destroy vpx codec");
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
pub struct DecodeFrames<'a> {
|
||
ctx: &'a mut vpx_codec_ctx_t,
|
||
iter: vpx_codec_iter_t,
|
||
}
|
||
|
||
impl<'a> Iterator for DecodeFrames<'a> {
|
||
type Item = Image;
|
||
fn next(&mut self) -> Option<Self::Item> {
|
||
let img = unsafe { vpx_codec_get_frame(self.ctx, &mut self.iter) };
|
||
if img.is_null() {
|
||
return None;
|
||
} else {
|
||
return Some(Image(img));
|
||
}
|
||
}
|
||
}
|
||
|
||
// https://chromium.googlesource.com/webm/libvpx/+/bali/vpx/src/vpx_image.c
|
||
pub struct Image(*mut vpx_image_t);
|
||
impl Image {
|
||
#[inline]
|
||
pub fn new() -> Self {
|
||
Self(std::ptr::null_mut())
|
||
}
|
||
|
||
#[inline]
|
||
pub fn is_null(&self) -> bool {
|
||
self.0.is_null()
|
||
}
|
||
|
||
#[inline]
|
||
pub fn width(&self) -> usize {
|
||
self.inner().d_w as _
|
||
}
|
||
|
||
#[inline]
|
||
pub fn height(&self) -> usize {
|
||
self.inner().d_h as _
|
||
}
|
||
|
||
#[inline]
|
||
pub fn format(&self) -> vpx_img_fmt_t {
|
||
// VPX_IMG_FMT_I420
|
||
self.inner().fmt
|
||
}
|
||
|
||
#[inline]
|
||
pub fn inner(&self) -> &vpx_image_t {
|
||
unsafe { &*self.0 }
|
||
}
|
||
|
||
#[inline]
|
||
pub fn stride(&self, iplane: usize) -> i32 {
|
||
self.inner().stride[iplane]
|
||
}
|
||
|
||
pub fn rgb(&self, stride_align: usize, rgba: bool, dst: &mut Vec<u8>) {
|
||
let h = self.height();
|
||
let mut w = self.width();
|
||
let bps = if rgba { 4 } else { 3 };
|
||
w = (w + stride_align - 1) & !(stride_align - 1);
|
||
dst.resize(h * w * bps, 0);
|
||
let img = self.inner();
|
||
unsafe {
|
||
if rgba {
|
||
super::I420ToARGB(
|
||
img.planes[0],
|
||
img.stride[0],
|
||
img.planes[1],
|
||
img.stride[1],
|
||
img.planes[2],
|
||
img.stride[2],
|
||
dst.as_mut_ptr(),
|
||
(w * bps) as _,
|
||
self.width() as _,
|
||
self.height() as _,
|
||
);
|
||
} else {
|
||
super::I420ToRAW(
|
||
img.planes[0],
|
||
img.stride[0],
|
||
img.planes[1],
|
||
img.stride[1],
|
||
img.planes[2],
|
||
img.stride[2],
|
||
dst.as_mut_ptr(),
|
||
(w * bps) as _,
|
||
self.width() as _,
|
||
self.height() as _,
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
#[inline]
|
||
pub fn data(&self) -> (&[u8], &[u8], &[u8]) {
|
||
unsafe {
|
||
let img = self.inner();
|
||
let h = (img.d_h as usize + 1) & !1;
|
||
let n = img.stride[0] as usize * h;
|
||
let y = slice::from_raw_parts(img.planes[0], n);
|
||
let n = img.stride[1] as usize * (h >> 1);
|
||
let u = slice::from_raw_parts(img.planes[1], n);
|
||
let v = slice::from_raw_parts(img.planes[2], n);
|
||
(y, u, v)
|
||
}
|
||
}
|
||
}
|
||
|
||
impl Drop for Image {
|
||
fn drop(&mut self) {
|
||
if !self.0.is_null() {
|
||
unsafe { vpx_img_free(self.0) };
|
||
}
|
||
}
|
||
}
|
||
|
||
unsafe impl Send for vpx_codec_ctx_t {}
|