From 70074749bad44c69fef1fb728ed24b656eaf8ed1 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Mon, 14 Apr 2025 20:26:16 +0200 Subject: [PATCH] adding avif support within rav1e --- node-hub/dora-dav1d/src/lib.rs | 13 ++ node-hub/dora-rav1e/Cargo.toml | 1 + node-hub/dora-rav1e/src/lib.rs | 375 +++++++++++++++++++-------------- node-hub/dora-rerun/src/lib.rs | 2 +- 4 files changed, 233 insertions(+), 158 deletions(-) diff --git a/node-hub/dora-dav1d/src/lib.rs b/node-hub/dora-dav1d/src/lib.rs index aad95ebe..a829f5aa 100644 --- a/node-hub/dora-dav1d/src/lib.rs +++ b/node-hub/dora-dav1d/src/lib.rs @@ -58,6 +58,19 @@ pub fn lib_main() -> Result<()> { }) => { if let Some(data) = data.as_any().downcast_ref::() { let data = data.values().clone(); + let encoding = metadata + .parameters + .get("encoding") + .and_then(|p| match p { + dora_node_api::Parameter::String(s) => Some(s), + _ => None, + }) + .map(|s| s.as_str()) + .unwrap_or("av1"); + if encoding != "av1" { + warn!("Unsupported encoding {}", encoding); + continue; + } match dec.send_data(data, None, None, None) { Err(e) => { warn!("Error sending data to the decoder: {}", e); diff --git a/node-hub/dora-rav1e/Cargo.toml b/node-hub/dora-rav1e/Cargo.toml index 109120b8..c2e35cd2 100644 --- a/node-hub/dora-rav1e/Cargo.toml +++ b/node-hub/dora-rav1e/Cargo.toml @@ -25,6 +25,7 @@ pyo3 = { workspace = true, features = [ "eyre", "generate-import-lib", ], optional = true } +avif-serialize = "0.8.3" [lib] diff --git a/node-hub/dora-rav1e/src/lib.rs b/node-hub/dora-rav1e/src/lib.rs index 44163f76..6884f7b6 100644 --- a/node-hub/dora-rav1e/src/lib.rs +++ b/node-hub/dora-rav1e/src/lib.rs @@ -10,14 +10,47 @@ use std::env::var; use dora_node_api::arrow::array::{UInt16Array, UInt8Array}; -use dora_node_api::{DoraNode, Event, IntoArrow, Parameter}; +use dora_node_api::dora_core::config::DataId; +use dora_node_api::{DoraNode, Event, IntoArrow, Metadata, Parameter}; use eyre::{Context as EyreContext, Result}; use log::warn; +use rav1e::color::{ColorDescription, MatrixCoefficients}; // Encode the same tiny blank frame 30 times use rav1e::config::SpeedSettings; use rav1e::*; +pub fn fill_zeros_toward_center_y_plane_in_place(y: &mut [u16], width: usize, height: usize) { + assert_eq!(y.len(), width * height); + + for row in 0..height { + let row_start = row * width; + let center = width / 2; + + // --- Fill left half (left to center) + let mut last = 0u16; + for col in 0..center { + let idx = row_start + col; + if y[idx] == 0 { + y[idx] = last; + } else { + last = y[idx]; + } + } + + // --- Fill right half (right to center) + last = 0u16; + for col in (center..width).rev() { + let idx = row_start + col; + if y[idx] == 0 { + y[idx] = last; + } else { + last = y[idx]; + } + } + } +} + fn bgr8_to_yuv420(bgr_data: Vec, width: usize, height: usize) -> (Vec, Vec, Vec) { let mut y_plane = vec![0; width * height]; let mut u_plane = vec![0; (width / 2) * (height / 2)]; @@ -69,6 +102,123 @@ fn get_yuv_planes(buffer: &[u8], width: usize, height: usize) -> (&[u8], &[u8], (y_plane, u_plane, v_plane) } +fn send_yuv( + y: &[u8], + u: &[u8], + v: &[u8], + enc: EncoderConfig, + width: usize, + height: usize, + node: &mut DoraNode, + id: DataId, + metadata: &mut Metadata, + output_enconding: &str, +) -> () { + // Create a new Arrow array for the YUV420 data + let cfg = Config::new().with_encoder_config(enc.clone()); + let mut ctx: Context = cfg.new_context().unwrap(); + let mut f = ctx.new_frame(); + + let xdec = f.planes[0].cfg.xdec; + let stride = (width + xdec) >> xdec; + f.planes[0].copy_from_raw_u8(&y, stride, 1); + let xdec = f.planes[1].cfg.xdec; + let stride = (width + xdec) >> xdec; + f.planes[1].copy_from_raw_u8(&u, stride, 1); + let xdec = f.planes[2].cfg.xdec; + let stride = (width + xdec) >> xdec; + f.planes[2].copy_from_raw_u8(&v, stride, 1); + + match ctx.send_frame(f) { + Ok(_) => {} + Err(e) => match e { + EncoderStatus::EnoughData => { + warn!("Unable to send frame "); + } + _ => { + warn!("Unable to send frame "); + } + }, + } + ctx.flush(); + match ctx.receive_packet() { + Ok(pkt) => match output_enconding { + "avif" => { + let data = pkt.data.clone(); + metadata.parameters.insert( + "encoding".to_string(), + Parameter::String("avif".to_string()), + ); + let matrix_coefficients = if let Some(desc) = enc.color_description { + desc.matrix_coefficients + } else { + MatrixCoefficients::BT709 + }; + let data = avif_serialize::Aviffy::new() + .set_chroma_subsampling((true, true)) + .set_seq_profile(0) + .matrix_coefficients(match matrix_coefficients { + MatrixCoefficients::Identity => { + avif_serialize::constants::MatrixCoefficients::Rgb + } + MatrixCoefficients::BT709 => { + avif_serialize::constants::MatrixCoefficients::Bt709 + } + MatrixCoefficients::Unspecified => { + avif_serialize::constants::MatrixCoefficients::Unspecified + } + MatrixCoefficients::BT601 => { + avif_serialize::constants::MatrixCoefficients::Bt601 + } + MatrixCoefficients::YCgCo => { + avif_serialize::constants::MatrixCoefficients::Ycgco + } + MatrixCoefficients::BT2020NCL => { + avif_serialize::constants::MatrixCoefficients::Bt2020Ncl + } + MatrixCoefficients::BT2020CL => { + avif_serialize::constants::MatrixCoefficients::Bt2020Cl + } + _ => { + warn!("color matrix coefficients"); + avif_serialize::constants::MatrixCoefficients::Rgb + } + }) + .to_vec( + &data, + None, + width as u32, + height as u32, + enc.bit_depth as u8, + ); + + let arrow = data.into_arrow(); + node.send_output(id, metadata.parameters.clone(), arrow) + .context("could not send output") + .unwrap(); + } + _ => { + metadata + .parameters + .insert("encoding".to_string(), Parameter::String("av1".to_string())); + let data = pkt.data; + let arrow = data.into_arrow(); + node.send_output(id, metadata.parameters.clone(), arrow) + .context("could not send output") + .unwrap(); + } + }, + Err(e) => match e { + EncoderStatus::LimitReached => {} + EncoderStatus::Encoded => {} + EncoderStatus::NeedMoreData => {} + _ => { + panic!("Unable to receive packet",); + } + }, + } +} + pub fn lib_main() -> Result<()> { let mut height = std::env::var("IMAGE_HEIGHT") .unwrap_or_else(|_| "480".to_string()) @@ -79,10 +229,17 @@ pub fn lib_main() -> Result<()> { .parse() .unwrap(); let speed = var("RAV1E_SPEED").map(|s| s.parse().unwrap()).unwrap_or(10); + let output_encoding = var("ENCODING").unwrap_or("av1".to_string()); let mut enc = EncoderConfig { width, height, speed_settings: SpeedSettings::from_preset(speed), + chroma_sampling: color::ChromaSampling::Cs420, + color_description: Some(ColorDescription { + matrix_coefficients: MatrixCoefficients::BT709, + transfer_characteristics: color::TransferCharacteristics::BT709, + color_primaries: color::ColorPrimaries::BT709, + }), low_latency: true, ..Default::default() }; @@ -127,115 +284,49 @@ pub fn lib_main() -> Result<()> { _ => {} } - let cfg = Config::new().with_encoder_config(enc.clone()); if encoding == "bgr8" { let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); let buffer: Vec = buffer.values().to_vec(); let (y, u, v) = bgr8_to_yuv420(buffer, width, height); - - // Transpose values from BGR to RGB - // let buffer: Vec = buffer.chunks(3).flat_map(|x| [x[2], x[1], x[0]]).collect(); - - let mut ctx: Context = cfg.new_context().unwrap(); - let mut f = ctx.new_frame(); - - let xdec = f.planes[0].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[0].copy_from_raw_u8(&y, stride, 1); - let xdec = f.planes[1].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[1].copy_from_raw_u8(&u, stride, 1); - let xdec = f.planes[2].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[2].copy_from_raw_u8(&v, stride, 1); - - match ctx.send_frame(f) { - Ok(_) => {} - Err(e) => match e { - EncoderStatus::EnoughData => { - warn!("Unable to send frame "); - continue; - } - _ => { - warn!("Unable to send frame "); - continue; - } - }, - } - metadata - .parameters - .insert("encoding".to_string(), Parameter::String("av1".to_string())); - ctx.flush(); - match ctx.receive_packet() { - Ok(pkt) => { - let data = pkt.data; - let arrow = data.into_arrow(); - node.send_output(id, metadata.parameters, arrow) - .context("could not send output") - .unwrap(); - } - Err(e) => match e { - EncoderStatus::LimitReached => {} - EncoderStatus::Encoded => {} - EncoderStatus::NeedMoreData => {} - _ => { - panic!("Unable to receive packet",); - } - }, - } + send_yuv( + &y, + &u, + &v, + enc, + width, + height, + &mut node, + id, + &mut metadata, + &output_encoding, + ); } else if encoding == "yuv420" { let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); let buffer = buffer.values(); //.to_vec(); let (y, u, v) = get_yuv_planes(buffer, width, height); - let mut ctx: Context = cfg.new_context().unwrap(); - let mut f = ctx.new_frame(); - - let xdec = f.planes[0].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[0].copy_from_raw_u8(&y, stride, 1); - let xdec = f.planes[1].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[1].copy_from_raw_u8(&u, stride, 1); - let xdec = f.planes[2].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[2].copy_from_raw_u8(&v, stride, 1); - - match ctx.send_frame(f) { - Ok(_) => {} - Err(e) => match e { - EncoderStatus::EnoughData => { - warn!("Unable to send frame "); - } - _ => { - warn!("Unable to send frame "); - } - }, - } - metadata - .parameters - .insert("encoding".to_string(), Parameter::String("av1".to_string())); - ctx.flush(); - match ctx.receive_packet() { - Ok(pkt) => { - let data = pkt.data; - let arrow = data.into_arrow(); - node.send_output(id, metadata.parameters, arrow) - .context("could not send output") - .unwrap(); - } - Err(e) => match e { - EncoderStatus::LimitReached => {} - EncoderStatus::Encoded => {} - EncoderStatus::NeedMoreData => {} - _ => { - panic!("Unable to receive packet",); - } - }, - } - } else if encoding == "mono16" { + send_yuv( + &y, + &u, + &v, + enc, + width, + height, + &mut node, + id, + &mut metadata, + &output_encoding, + ); + } else if encoding == "mono16" || encoding == "z16" { let buffer: &UInt16Array = data.as_any().downcast_ref().unwrap(); - let buffer: &[u16] = buffer.values(); + let mut buffer = buffer.values().to_vec(); + if std::env::var("FILL_ZEROS") + .map(|s| s != "false") + .unwrap_or(true) + { + fill_zeros_toward_center_y_plane_in_place(&mut buffer, width, height); + } + // let buffer = shift_u16_slice_to_upper_12_bits(buffer); let bytes: &[u8] = &bytemuck::cast_slice(&buffer); @@ -258,17 +349,25 @@ pub fn lib_main() -> Result<()> { } }, } - metadata - .parameters - .insert("encoding".to_string(), Parameter::String("av1".to_string())); ctx.flush(); match ctx.receive_packet() { Ok(pkt) => { let data = pkt.data; - let arrow = data.into_arrow(); - node.send_output(id, metadata.parameters, arrow) - .context("could not send output") - .unwrap(); + match output_encoding.as_str() { + "avif" => { + warn!("avif encoding not supported for mono16"); + } + _ => { + metadata.parameters.insert( + "encoding".to_string(), + Parameter::String("av1".to_string()), + ); + let arrow = data.into_arrow(); + node.send_output(id, metadata.parameters, arrow) + .context("could not send output") + .unwrap(); + } + } } Err(e) => match e { EncoderStatus::LimitReached => {} @@ -285,56 +384,18 @@ pub fn lib_main() -> Result<()> { let buffer: Vec = buffer.chunks(3).flat_map(|x| [x[2], x[1], x[0]]).collect(); let (y, u, v) = bgr8_to_yuv420(buffer, width, height); - - // Transpose values from BGR to RGB - - let mut ctx: Context = cfg.new_context().unwrap(); - let mut f = ctx.new_frame(); - - let xdec = f.planes[0].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[0].copy_from_raw_u8(&y, stride, 1); - let xdec = f.planes[1].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[1].copy_from_raw_u8(&u, stride, 1); - let xdec = f.planes[2].cfg.xdec; - let stride = (width + xdec) >> xdec; - f.planes[2].copy_from_raw_u8(&v, stride, 1); - - match ctx.send_frame(f) { - Ok(_) => {} - Err(e) => match e { - EncoderStatus::EnoughData => { - warn!("Unable to send frame "); - continue; - } - _ => { - warn!("Unable to send frame "); - continue; - } - }, - } - metadata - .parameters - .insert("encoding".to_string(), Parameter::String("av1".to_string())); - ctx.flush(); - match ctx.receive_packet() { - Ok(pkt) => { - let data = pkt.data; - let arrow = data.into_arrow(); - node.send_output(id, metadata.parameters, arrow) - .context("could not send output") - .unwrap(); - } - Err(e) => match e { - EncoderStatus::LimitReached => {} - EncoderStatus::Encoded => {} - EncoderStatus::NeedMoreData => {} - _ => { - panic!("Unable to receive packet",); - } - }, - } + send_yuv( + &y, + &u, + &v, + enc, + width, + height, + &mut node, + id, + &mut metadata, + &output_encoding, + ); } else { unimplemented!("We haven't worked on additional encodings."); } diff --git a/node-hub/dora-rerun/src/lib.rs b/node-hub/dora-rerun/src/lib.rs index c6884630..2da325ac 100644 --- a/node-hub/dora-rerun/src/lib.rs +++ b/node-hub/dora-rerun/src/lib.rs @@ -155,7 +155,7 @@ pub fn lib_main() -> Result<()> { ); rec.log(id.as_str(), &image) .context("could not log image")?; - } else if ["jpeg", "png"].contains(&encoding) { + } else if ["jpeg", "png", "avif"].contains(&encoding) { let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); let buffer: &[u8] = buffer.values();