use bytes::Bytes;
use super::backend::{self, Backend};
use crate::frame::{Frame, I420};
use crate::{Error, Size};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum Codec {
#[default]
H264,
H265,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum Kind {
#[default]
Auto,
Hardware,
Software,
Named(String),
}
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct Config {
pub width: u32,
pub height: u32,
pub framerate: u32,
pub bitrate: Option<u64>,
pub gop: u32,
pub codec: Codec,
pub kind: Kind,
}
impl Config {
pub fn new(width: u32, height: u32, framerate: u32) -> Self {
Self {
width,
height,
framerate,
bitrate: None,
gop: framerate.saturating_mul(2).max(1),
codec: Codec::default(),
kind: Kind::Auto,
}
}
pub fn size(&self) -> Size {
Size::new(self.width, self.height)
}
pub(crate) fn resolved_bitrate(&self) -> u64 {
self.bitrate.unwrap_or_else(|| {
((self.size().pixels() * self.framerate as u64) as f64 * 0.07) as u64
})
}
}
pub struct Encoder {
backend: Box<dyn Backend>,
codec: Codec,
size: Size,
bitrate: u64,
}
impl Encoder {
pub fn new(config: &Config) -> Result<Self, Error> {
if config.framerate == 0 {
return Err(Error::InvalidFramerate(0));
}
let size = config.size();
size.validate("encoder")?;
let backend = backend::open(config)?;
Ok(Self {
backend,
codec: config.codec,
size,
bitrate: config.resolved_bitrate(),
})
}
pub fn name(&self) -> &str {
self.backend.name()
}
pub fn size(&self) -> Size {
self.size
}
pub fn bitrate(&self) -> u64 {
self.bitrate
}
pub fn set_bitrate(&mut self, bitrate: u64) -> Result<(), Error> {
if bitrate == self.bitrate {
return Ok(());
}
self.backend.set_bitrate(bitrate)?;
self.bitrate = bitrate;
Ok(())
}
pub fn codec(&self) -> Codec {
self.codec
}
pub fn encode_rgba(&mut self, rgba: &[u8], size: Size, keyframe: bool) -> Result<Vec<Bytes>, Error> {
self.check_frame(size, rgba.len(), size.pixels() as usize * 4, "RGBA")?;
let frame = Frame::I420(I420::from_rgba(rgba, size.width * 4, size.width, size.height)?);
self.encode_raw(&frame, keyframe)
}
pub fn encode_i420(&mut self, i420: &[u8], size: Size, keyframe: bool) -> Result<Vec<Bytes>, Error> {
self.check_frame(size, i420.len(), I420::len(size.width, size.height), "I420")?;
let frame = Frame::I420(I420 {
width: size.width,
height: size.height,
data: i420.to_vec(),
});
self.encode_raw(&frame, keyframe)
}
fn check_frame(&self, size: Size, got: usize, expected: usize, what: &str) -> Result<(), Error> {
if size != self.size {
return Err(Error::Codec(anyhow::anyhow!(
"frame {size} does not match encoder {}",
self.size
)));
}
if got != expected {
return Err(Error::Codec(anyhow::anyhow!(
"{what} buffer is {got} bytes, expected {expected} for {size}"
)));
}
Ok(())
}
pub fn encode(&mut self, frame: &crate::decode::Frame, keyframe: bool) -> Result<Vec<Bytes>, Error> {
self.encode_raw(&frame.inner, keyframe)
}
pub(crate) fn encode_raw(&mut self, frame: &Frame, keyframe: bool) -> Result<Vec<Bytes>, Error> {
let size = Size::new(frame.width(), frame.height());
if size != self.size {
return Err(Error::Codec(anyhow::anyhow!(
"frame {size} does not match encoder {}",
self.size
)));
}
self.backend.encode(frame, keyframe)
}
pub fn finish(mut self) -> Result<Vec<Bytes>, Error> {
self.backend.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn gray_rgba(width: u32, height: u32) -> Vec<u8> {
vec![0x80u8; width as usize * height as usize * 4]
}
#[test]
fn software_encoder_emits_annexb() {
let config = Config {
kind: Kind::Software,
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).expect("openh264 is vendored, always available");
assert_eq!(encoder.name(), "openh264");
let frame = gray_rgba(320, 240);
let mut packets = Vec::new();
for i in 0..30 {
packets.extend(encoder.encode_rgba(&frame, Size::new(320, 240), i == 0).unwrap());
}
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty(), "encoder produced no packets");
let first = &packets[0];
let has_start_code = first.starts_with(&[0, 0, 0, 1]) || first.starts_with(&[0, 0, 1]);
assert!(
has_start_code,
"first packet is not Annex-B: {:02x?}",
&first[..first.len().min(8)]
);
}
#[test]
fn encode_rgba_emits_annexb() {
let config = Config {
kind: Kind::Software,
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).unwrap();
let rgba = gray_rgba(320, 240);
let mut packets = encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty());
assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
}
#[test]
fn encode_i420_emits_annexb() {
let config = Config {
kind: Kind::Software,
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).unwrap();
let data = vec![0x80u8; I420::len(320, 240)];
let mut packets = encoder.encode_i420(&data, Size::new(320, 240), true).unwrap();
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty());
assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
}
#[test]
fn encode_i420_rejects_wrong_size() {
let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
return;
};
assert!(matches!(
encoder.encode_i420(&[0u8; 16], Size::new(320, 240), false),
Err(Error::Codec(_))
));
}
#[test]
fn encode_rgba_rejects_short_buffer() {
let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
return;
};
assert!(matches!(
encoder.encode_rgba(&[0u8; 16], Size::new(320, 240), false),
Err(Error::Codec(_))
));
}
#[test]
fn encode_rgba_rejects_dimension_mismatch() {
let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
return;
};
let rgba = gray_rgba(640, 480);
assert!(matches!(
encoder.encode_rgba(&rgba, Size::new(640, 480), false),
Err(Error::Codec(_))
));
}
#[test]
fn encode_i420_rejects_oversized_buffer() {
let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
return;
};
let data = vec![0x80u8; I420::len(640, 480)];
assert!(matches!(
encoder.encode_i420(&data, Size::new(640, 480), false),
Err(Error::Codec(_))
));
}
#[test]
fn encode_rejects_transposed_frame() {
let Ok(mut encoder) = Encoder::new(&Config::new(320, 240, 30)) else {
return;
};
let rgba = gray_rgba(240, 320);
assert_eq!(rgba.len(), gray_rgba(320, 240).len(), "the byte counts must collide");
assert!(matches!(
encoder.encode_rgba(&rgba, Size::new(240, 320), false),
Err(Error::Codec(_))
));
let i420 = vec![0x80u8; I420::len(240, 320)];
assert_eq!(i420.len(), I420::len(320, 240), "the byte counts must collide");
assert!(matches!(
encoder.encode_i420(&i420, Size::new(240, 320), false),
Err(Error::Codec(_))
));
}
#[test]
fn new_rejects_zero_framerate() {
let config = Config::new(320, 240, 0);
assert!(matches!(Encoder::new(&config), Err(Error::InvalidFramerate(0))));
}
#[test]
fn unknown_named_encoder_errors() {
let config = Config {
kind: Kind::Named("definitely_not_a_codec".into()),
..Config::new(320, 240, 30)
};
assert!(matches!(Encoder::new(&config), Err(Error::NoEncoder(_))));
}
#[cfg(target_os = "macos")]
#[test]
fn videotoolbox_emits_annexb_keyframe() {
let config = Config {
kind: Kind::Named("videotoolbox".into()),
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).expect("videotoolbox is available on macOS");
assert_eq!(encoder.name(), "videotoolbox");
let frame = gray_rgba(320, 240);
let mut packets = Vec::new();
for i in 0..10 {
packets.extend(encoder.encode_rgba(&frame, Size::new(320, 240), i == 0).unwrap());
}
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty(), "encoder produced no packets");
let first = &packets[0];
assert!(
first.starts_with(&[0, 0, 0, 1]) || first.starts_with(&[0, 0, 1]),
"first packet is not Annex-B"
);
let types = nal_types(first);
assert!(types.contains(&7), "no SPS in first packet: {types:?}");
assert!(types.contains(&8), "no PPS in first packet: {types:?}");
assert!(types.contains(&5), "first packet is not an IDR: {types:?}");
}
#[cfg(target_os = "macos")]
#[test]
fn videotoolbox_emits_annexb_keyframe_h265() {
let config = Config {
codec: Codec::H265,
kind: Kind::Named("videotoolbox".into()),
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).expect("videotoolbox HEVC is available on macOS");
assert_eq!(encoder.name(), "videotoolbox");
assert_eq!(encoder.codec(), Codec::H265);
let frame = gray_rgba(320, 240);
let mut packets = Vec::new();
for i in 0..10 {
packets.extend(encoder.encode_rgba(&frame, Size::new(320, 240), i == 0).unwrap());
}
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty(), "encoder produced no packets");
let first = &packets[0];
assert!(
first.starts_with(&[0, 0, 0, 1]) || first.starts_with(&[0, 0, 1]),
"first packet is not Annex-B"
);
let types = hevc_nal_types(first);
assert!(types.contains(&32), "no VPS in first packet: {types:?}");
assert!(types.contains(&33), "no SPS in first packet: {types:?}");
assert!(types.contains(&34), "no PPS in first packet: {types:?}");
assert!(
types.iter().any(|t| (16..=23).contains(t)),
"first packet is not an IRAP: {types:?}"
);
}
#[cfg(target_os = "macos")]
fn hevc_nal_types(annexb: &[u8]) -> Vec<u8> {
let mut types = Vec::new();
let mut i = 0;
while i + 3 < annexb.len() {
if annexb[i..i + 3] == [0, 0, 1] {
types.push((annexb[i + 3] >> 1) & 0x3f);
i += 3;
} else {
i += 1;
}
}
types
}
#[cfg(target_os = "macos")]
#[test]
fn videotoolbox_encodes_surface_zero_copy() {
let config = Config {
kind: Kind::Named("videotoolbox".into()),
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).unwrap();
let mut packets = Vec::new();
for i in 0..10 {
let frame = Frame::Surface(nv12_surface(320, 240));
packets.extend(encoder.encode_raw(&frame, i == 0).unwrap());
}
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty());
let types = nal_types(&packets[0]);
assert!(
types.contains(&7) && types.contains(&8) && types.contains(&5),
"no IDR: {types:?}"
);
}
#[cfg(target_os = "macos")]
#[test]
fn openh264_downloads_surface() {
let config = Config {
kind: Kind::Software,
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).unwrap();
let frame = Frame::Surface(nv12_surface(320, 240));
let mut packets = encoder.encode_raw(&frame, true).unwrap();
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty());
assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
}
#[cfg(target_os = "macos")]
fn nv12_surface(width: u32, height: u32) -> crate::frame::macos::Surface {
use std::ptr::{self, NonNull};
use objc2_core_foundation::CFRetained;
use objc2_core_video::{
CVPixelBuffer, CVPixelBufferCreate, CVPixelBufferGetBaseAddressOfPlane, CVPixelBufferGetBytesPerRowOfPlane,
CVPixelBufferLockBaseAddress, CVPixelBufferLockFlags, CVPixelBufferUnlockBaseAddress,
kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange,
};
let mut raw: *mut CVPixelBuffer = ptr::null_mut();
let status = unsafe {
CVPixelBufferCreate(
None,
width as usize,
height as usize,
kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange,
None,
NonNull::new(&mut raw).unwrap(),
)
};
assert_eq!(status, 0, "CVPixelBufferCreate failed");
let buffer = unsafe { CFRetained::from_raw(NonNull::new(raw).unwrap()) };
let flags = CVPixelBufferLockFlags(0);
assert_eq!(unsafe { CVPixelBufferLockBaseAddress(&buffer, flags) }, 0);
for (plane, rows) in [(0usize, height as usize), (1usize, height as usize / 2)] {
let base = CVPixelBufferGetBaseAddressOfPlane(&buffer, plane) as *mut u8;
let stride = CVPixelBufferGetBytesPerRowOfPlane(&buffer, plane);
unsafe { ptr::write_bytes(base, 128, stride * rows) };
}
unsafe { CVPixelBufferUnlockBaseAddress(&buffer, flags) };
crate::frame::macos::Surface::new(buffer, width, height)
}
#[cfg(any(target_os = "macos", target_os = "windows"))]
fn nal_types(annexb: &[u8]) -> Vec<u8> {
let mut types = Vec::new();
let mut i = 0;
while i + 3 < annexb.len() {
if annexb[i..i + 3] == [0, 0, 1] {
types.push(annexb[i + 3] & 0x1f);
i += 3;
} else {
i += 1;
}
}
types
}
#[cfg(target_os = "windows")]
#[test]
#[ignore]
fn mediafoundation_cpu_rgba() {
let config = Config {
kind: Kind::Named("mediafoundation".into()),
..Config::new(640, 480, 30)
};
let mut encoder = Encoder::new(&config).expect("hardware H.264 encoder available");
assert_eq!(encoder.name(), "mediafoundation");
let frame = gray_rgba(640, 480);
let mut packets = Vec::new();
for i in 0..30 {
packets.extend(encoder.encode_rgba(&frame, Size::new(640, 480), i == 0).unwrap());
}
packets.extend(encoder.finish().unwrap());
assert!(!packets.is_empty(), "encoder produced no packets");
let types = nal_types(&packets[0]);
assert!(types.contains(&7), "no SPS in first packet: {types:?}");
assert!(types.contains(&8), "no PPS in first packet: {types:?}");
assert!(types.contains(&5), "first packet is not an IDR: {types:?}");
}
#[cfg(target_os = "windows")]
#[tokio::test]
#[ignore]
async fn mediafoundation_camera_texture() {
let mut camera = crate::capture::open(&crate::capture::Config::default())
.await
.expect("open default camera");
let (w, h) = (camera.width(), camera.height());
let config = Config {
kind: Kind::Named("mediafoundation".into()),
..Config::new(w, h, camera.framerate().unwrap_or(30))
};
let mut encoder = Encoder::new(&config).expect("hardware H.264 encoder available");
let mut packets = Vec::new();
let mut textures = 0;
for i in 0..30 {
let frame = camera.read().await.expect("frame, not end of stream");
if matches!(frame, Frame::Texture(_)) {
textures += 1;
}
packets.extend(encoder.encode_raw(&frame, i == 0).unwrap());
}
packets.extend(encoder.finish().unwrap());
assert!(textures > 0, "capture never produced a GPU texture");
assert!(!packets.is_empty(), "encoder produced no packets");
let types = nal_types(&packets[0]);
assert!(
types.contains(&7) && types.contains(&8) && types.contains(&5),
"no IDR: {types:?}"
);
}
#[test]
fn set_bitrate_retunes_software_encoder() {
let config = Config {
kind: Kind::Software,
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).unwrap();
let rgba = gray_rgba(320, 240);
let opened = encoder.bitrate();
assert_eq!(opened, config.resolved_bitrate());
encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
let halved = opened / 2;
encoder.set_bitrate(halved).unwrap();
assert_eq!(encoder.bitrate(), halved);
let packets = encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
assert!(!packets.is_empty(), "encoder produced nothing after a retune");
assert!(packets[0].starts_with(&[0, 0, 0, 1]) || packets[0].starts_with(&[0, 0, 1]));
}
#[test]
fn set_bitrate_before_the_first_frame_is_deferred() {
let config = Config {
kind: Kind::Software,
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).unwrap();
let halved = encoder.bitrate() / 2;
encoder.set_bitrate(halved).expect("a retune before the first frame");
assert_eq!(encoder.bitrate(), halved);
let rgba = gray_rgba(320, 240);
let packets = encoder.encode_rgba(&rgba, Size::new(320, 240), true).unwrap();
assert!(!packets.is_empty());
encoder.set_bitrate(halved / 2).unwrap();
assert!(encoder.encode_rgba(&rgba, Size::new(320, 240), false).is_ok());
}
#[test]
fn set_bitrate_to_current_is_a_noop() {
let config = Config {
kind: Kind::Software,
..Config::new(320, 240, 30)
};
let mut encoder = Encoder::new(&config).unwrap();
let opened = encoder.bitrate();
encoder.set_bitrate(opened).unwrap();
assert_eq!(encoder.bitrate(), opened);
}
#[test]
fn default_bitrate_scales_with_resolution() {
let small = Config::new(320, 240, 30).resolved_bitrate();
let large = Config::new(1920, 1080, 30).resolved_bitrate();
assert!(large > small);
assert!(small > 0);
}
}