#[cfg_attr(not(all(target_os = "linux", feature = "v4l2")), allow(dead_code))] pub(crate) mod h264;
#[cfg(all(target_os = "linux", feature = "v4l2"))]
mod v4l2;
use std::sync::{Arc, Mutex};
use edgefirst_tensor::{Colorimetry, PixelFormat, Region, TensorDyn};
pub use h264::{access_units, AccessUnits};
use crate::{CodecError, Result};
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum VideoCodec {
H264,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Bitrate {
#[default]
Auto,
Bps(u32),
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum H264Profile {
Baseline,
ConstrainedBaseline,
Main,
High,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[allow(missing_docs)]
pub enum H264Level {
L1_0,
L1b,
L1_1,
L1_2,
L1_3,
L2_0,
L2_1,
L2_2,
L3_0,
L3_1,
L3_2,
L4_0,
L4_1,
L4_2,
L5_0,
L5_1,
L5_2,
L6_0,
L6_1,
L6_2,
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq)]
pub struct EncoderConfig {
pub codec: VideoCodec,
pub width: u32,
pub height: u32,
pub input_format: PixelFormat,
pub frame_rate: f64,
pub bitrate: Bitrate,
pub gop: Option<u32>,
pub repeat_headers: bool,
pub profile: Option<H264Profile>,
pub level: Option<H264Level>,
}
impl EncoderConfig {
pub fn h264(width: u32, height: u32, input_format: PixelFormat, frame_rate: f64) -> Self {
Self {
codec: VideoCodec::H264,
width,
height,
input_format,
frame_rate,
bitrate: Bitrate::Auto,
gop: None,
repeat_headers: true,
profile: None,
level: None,
}
}
pub fn validate(&self) -> Result<()> {
let invalid = |what: String| Err(CodecError::InvalidConfig(what));
if self.width == 0 || self.height == 0 {
return invalid(format!("size {}x{} is empty", self.width, self.height));
}
let (hsub, vsub) = chroma_subsampling(self.input_format);
if !self.width.is_multiple_of(hsub) || !self.height.is_multiple_of(vsub) {
return invalid(format!(
"size {}x{} is not a whole number of {} chroma blocks",
self.width, self.height, self.input_format
));
}
if !self.frame_rate.is_finite() || self.frame_rate <= 0.0 {
return invalid(format!("frame rate {} is not positive", self.frame_rate));
}
if self.bitrate == Bitrate::Bps(0) {
return invalid("bitrate is 0".into());
}
if self.gop == Some(0) {
return invalid("GOP is 0".into());
}
Ok(())
}
}
pub(crate) fn chroma_subsampling(format: PixelFormat) -> (u32, u32) {
match format {
PixelFormat::Nv12 => (2, 2),
PixelFormat::Nv16 | PixelFormat::Yuyv | PixelFormat::Vyuy => (2, 1),
_ => (1, 1),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct FrameOptions {
pub pts: u64,
pub crop: Option<Region>,
pub keyframe: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EncodedFrame {
pub data: Vec<u8>,
pub keyframe: bool,
pub pts: u64,
}
pub(crate) trait EncoderBackend: Send {
fn encode(&mut self, src: &TensorDyn, opts: &FrameOptions) -> Result<Option<EncodedFrame>>;
fn flush(&mut self) -> Result<Vec<EncodedFrame>>;
fn device(&self) -> String;
}
pub struct VideoEncoder {
backend: Box<dyn EncoderBackend>,
}
impl std::fmt::Debug for VideoEncoder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("VideoEncoder")
.field("device", &self.backend.device())
.finish()
}
}
impl VideoEncoder {
pub fn new(config: EncoderConfig) -> Result<Self> {
config.validate()?;
#[cfg(all(target_os = "linux", feature = "v4l2"))]
{
let backend = v4l2::encoder::V4l2Encoder::open(config)?;
Ok(Self {
backend: Box::new(backend),
})
}
#[cfg(not(all(target_os = "linux", feature = "v4l2")))]
{
let _ = config;
Err(CodecError::NoDevice(
"no video encoder backend in this build".into(),
))
}
}
pub fn encode(&mut self, src: &TensorDyn, opts: &FrameOptions) -> Result<Option<EncodedFrame>> {
self.backend.encode(src, opts)
}
pub fn flush(&mut self) -> Result<Vec<EncodedFrame>> {
self.backend.flush()
}
pub fn device(&self) -> String {
self.backend.device()
}
}
#[derive(Debug)]
pub enum OutputPool {
Allocate {
count: usize,
},
Provided(Vec<TensorDyn>),
}
impl Default for OutputPool {
fn default() -> Self {
Self::Allocate { count: 4 }
}
}
#[non_exhaustive]
#[derive(Debug)]
pub struct DecoderConfig {
pub codec: VideoCodec,
pub output: OutputPool,
}
impl DecoderConfig {
pub fn h264() -> Self {
Self {
codec: VideoCodec::H264,
output: OutputPool::default(),
}
}
pub fn with_output(mut self, output: OutputPool) -> Self {
self.output = output;
self
}
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq)]
pub struct StreamInfo {
pub width: usize,
pub height: usize,
pub crop: Region,
pub coded_width: usize,
pub coded_height: usize,
pub format: PixelFormat,
pub colorimetry: Option<Colorimetry>,
pub buffers: usize,
pub memory: OutputMemory,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutputMemory {
Imported,
Exported,
}
#[derive(Debug)]
pub enum DecodeEvent {
Frame(DecodedFrame),
Reconfigured(StreamInfo),
}
#[derive(Debug, Default)]
pub(crate) struct Returns(Mutex<Vec<(u64, u32)>>);
impl Returns {
#[cfg_attr(not(all(target_os = "linux", feature = "v4l2")), allow(dead_code))] pub(crate) fn take(&self, generation: u64) -> Vec<u32> {
let mut returned = self.0.lock().unwrap_or_else(|e| e.into_inner());
returned
.drain(..)
.filter(|&(g, _)| g == generation)
.map(|(_, index)| index)
.collect()
}
}
#[derive(Debug)]
pub struct DecodedFrame {
tensor: Arc<TensorDyn>,
pts: u64,
corrupted: bool,
slot: (u64, u32),
returns: Arc<Returns>,
}
impl DecodedFrame {
#[cfg_attr(not(all(target_os = "linux", feature = "v4l2")), allow(dead_code))] pub(crate) fn new(
tensor: Arc<TensorDyn>,
pts: u64,
corrupted: bool,
slot: (u64, u32),
returns: Arc<Returns>,
) -> Self {
Self {
tensor,
pts,
corrupted,
slot,
returns,
}
}
pub fn tensor(&self) -> &TensorDyn {
&self.tensor
}
pub fn pts(&self) -> u64 {
self.pts
}
pub fn corrupted(&self) -> bool {
self.corrupted
}
}
impl Drop for DecodedFrame {
fn drop(&mut self) {
self.returns
.0
.lock()
.unwrap_or_else(|e| e.into_inner())
.push(self.slot);
}
}
pub(crate) trait DecoderBackend: Send {
fn decode(&mut self, access_unit: &[u8], pts: u64) -> Result<Vec<DecodeEvent>>;
fn flush(&mut self) -> Result<Vec<DecodeEvent>>;
fn stream_info(&self) -> Option<&StreamInfo>;
fn device(&self) -> String;
}
pub struct VideoDecoder {
backend: Box<dyn DecoderBackend>,
}
impl std::fmt::Debug for VideoDecoder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("VideoDecoder")
.field("device", &self.backend.device())
.field("stream", &self.backend.stream_info())
.finish()
}
}
impl VideoDecoder {
pub fn new(config: DecoderConfig) -> Result<Self> {
if let OutputPool::Provided(pool) = &config.output {
if pool.is_empty() {
return Err(CodecError::InvalidConfig(
"the provided pool is empty".into(),
));
}
if let Some(t) = pool
.iter()
.find(|t| t.memory() != edgefirst_tensor::TensorMemory::DmaBuf)
{
return Err(CodecError::InvalidConfig(format!(
"the provided pool has a tensor in {:?} memory; the decoder writes DMA-BUF tensors",
t.memory()
)));
}
}
#[cfg(all(target_os = "linux", feature = "v4l2"))]
{
let backend = v4l2::decoder::V4l2Decoder::open(config)?;
Ok(Self {
backend: Box::new(backend),
})
}
#[cfg(not(all(target_os = "linux", feature = "v4l2")))]
{
let _ = config;
Err(CodecError::NoDevice(
"no video decoder backend in this build".into(),
))
}
}
pub fn decode(&mut self, access_unit: &[u8], pts: u64) -> Result<Vec<DecodeEvent>> {
self.backend.decode(access_unit, pts)
}
pub fn flush(&mut self) -> Result<Vec<DecodeEvent>> {
self.backend.flush()
}
pub fn stream_info(&self) -> Option<&StreamInfo> {
self.backend.stream_info()
}
pub fn device(&self) -> String {
self.backend.device()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn nv12_1080p() -> EncoderConfig {
EncoderConfig::h264(1920, 1080, PixelFormat::Nv12, 30.0)
}
#[test]
fn default_h264_config_is_valid() {
let c = nv12_1080p();
assert!(c.validate().is_ok());
assert_eq!(c.codec, VideoCodec::H264);
assert_eq!(c.bitrate, Bitrate::Auto);
assert!(c.repeat_headers);
assert!(c.gop.is_none() && c.profile.is_none() && c.level.is_none());
}
#[test]
fn rejects_empty_and_odd_sizes() {
let mut c = nv12_1080p();
c.width = 0;
assert!(matches!(c.validate(), Err(CodecError::InvalidConfig(_))));
let mut c = nv12_1080p();
c.height = 1081;
assert!(matches!(c.validate(), Err(CodecError::InvalidConfig(_))));
let mut c = EncoderConfig::h264(641, 480, PixelFormat::Yuyv, 30.0);
assert!(matches!(c.validate(), Err(CodecError::InvalidConfig(_))));
c.width = 640;
c.height = 481;
assert!(c.validate().is_ok(), "YUYV subsamples horizontally only");
}
#[test]
fn rejects_bad_rate_bitrate_and_gop() {
for fps in [0.0, -1.0, f64::NAN, f64::INFINITY] {
let mut c = nv12_1080p();
c.frame_rate = fps;
assert!(c.validate().is_err(), "{fps}");
}
let mut c = nv12_1080p();
c.bitrate = Bitrate::Bps(0);
assert!(c.validate().is_err());
let mut c = nv12_1080p();
c.gop = Some(0);
assert!(c.validate().is_err());
c.gop = Some(1);
c.bitrate = Bitrate::Bps(4_000_000);
assert!(c.validate().is_ok());
}
#[test]
fn frames_can_be_shared_between_threads() {
fn shareable<T: Send + Sync>() {}
shareable::<DecodedFrame>();
shareable::<Arc<DecodedFrame>>();
}
#[test]
fn a_dropped_frame_returns_its_buffer_to_its_own_pool() {
let returns = Arc::new(Returns::default());
let tensor =
|n| Arc::new(TensorDyn::new(&[n], edgefirst_tensor::DType::U8, None, None).unwrap());
let old = DecodedFrame::new(tensor(16), 7, false, (1, 3), Arc::clone(&returns));
let current = DecodedFrame::new(tensor(16), 8, true, (2, 5), Arc::clone(&returns));
assert_eq!((current.pts(), current.corrupted()), (8, true));
assert!(
returns.take(2).is_empty(),
"nothing returned while frames are held"
);
drop(old);
drop(current);
assert_eq!(
returns.take(2),
vec![5],
"a buffer of an earlier pool is not returned"
);
assert!(returns.take(1).is_empty(), "take drains every generation");
}
#[test]
fn levels_are_ordered() {
assert!(H264Level::L1b < H264Level::L1_1);
assert!(H264Level::L4_0 < H264Level::L5_1);
}
}