From 70074749bad44c69fef1fb728ed24b656eaf8ed1 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Mon, 14 Apr 2025 20:26:16 +0200 Subject: [PATCH 1/4] 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(); From 48f72dce5a899204401a85f2e9c120cf39465e31 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Mon, 14 Apr 2025 23:18:00 +0200 Subject: [PATCH 2/4] Minor fix on rav1e dav1d --- Cargo.lock | 16 ++- examples/av1-encoding/dataflow.yml | 7 -- examples/av1-encoding/ios-dev.yaml | 53 +++------ node-hub/dora-dav1d/src/lib.rs | 105 +++++++++++++++--- .../dora-ios-lidar/dora_ios_lidar/main.py | 10 +- node-hub/dora-rav1e/src/lib.rs | 15 ++- node-hub/opencv-plot/opencv_plot/main.py | 18 ++- node-hub/opencv-plot/pyproject.toml | 8 +- 8 files changed, 168 insertions(+), 64 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e62a3a37..29727563 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1350,6 +1350,15 @@ dependencies = [ "v_frame", ] +[[package]] +name = "avif-serialize" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98922d6a4cfbcb08820c69d8eeccc05bb1f29bfa06b4f5b1dbfe9a868bd7608e" +dependencies = [ + "arrayvec", +] + [[package]] name = "axum" version = "0.7.9" @@ -3438,6 +3447,7 @@ dependencies = [ name = "dora-rav1e" version = "0.3.11+fix1" dependencies = [ + "avif-serialize", "bytemuck", "dora-node-api", "eyre", @@ -5405,7 +5415,7 @@ dependencies = [ "httpdate", "itoa", "pin-project-lite", - "socket2 0.4.10", + "socket2 0.5.8", "tokio", "tower-service", "tracing", @@ -6368,7 +6378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc2f4eb4bc735547cfed7c0a4922cbd04a4655978c09b54f1f7b228750664c34" dependencies = [ "cfg-if 1.0.0", - "windows-targets 0.48.5", + "windows-targets 0.52.6", ] [[package]] @@ -13923,7 +13933,7 @@ version = "1.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97fee6b57c6a41524a810daee9286c02d7752c4253064d0b05472833a438f675" dependencies = [ - "cfg-if 0.1.10", + "cfg-if 1.0.0", "static_assertions", ] diff --git a/examples/av1-encoding/dataflow.yml b/examples/av1-encoding/dataflow.yml index 131da27a..2dbcef01 100644 --- a/examples/av1-encoding/dataflow.yml +++ b/examples/av1-encoding/dataflow.yml @@ -42,9 +42,6 @@ nodes: image: dav1d-remote/image outputs: - image - env: - IMAGE_WIDTH: 640 - IMAGE_HEIGHT: 480 - id: dav1d-local path: dora-dav1d @@ -55,9 +52,6 @@ nodes: image: rav1e-remote/image outputs: - image - env: - IMAGE_WIDTH: 640 - IMAGE_HEIGHT: 480 - id: plot build: pip install -e ../../node-hub/dora-rerun @@ -66,4 +60,3 @@ nodes: path: dora-rerun inputs: image_decode: dav1d-local/image - image_echo: echo/image diff --git a/examples/av1-encoding/ios-dev.yaml b/examples/av1-encoding/ios-dev.yaml index d14fdb36..02fdf900 100644 --- a/examples/av1-encoding/ios-dev.yaml +++ b/examples/av1-encoding/ios-dev.yaml @@ -2,72 +2,53 @@ nodes: - id: camera build: pip install -e ../../node-hub/dora-ios-lidar path: dora-ios-lidar - _unstable_deploy: - machine: encoder-ios inputs: tick: dora/timer/millis/20 outputs: - image - depth env: - IMAGE_WIDTH: 640 - IMAGE_HEIGHT: 480 + IMAGE_WIDTH: 1280 + IMAGE_HEIGHT: 720 + ROTATE: ROTATE_90_CLOCKWISE - id: rav1e-local path: dora-rav1e build: cargo build -p dora-rav1e --release - _unstable_deploy: - machine: encoder-ios inputs: image: camera/image - depth: camera/depth outputs: - image - - depth env: - RAV1E_SPEED: 4 + RAV1E_SPEED: 10 - - id: dav1d-remote - path: dora-dav1d - build: cargo build -p dora-dav1d --release - _unstable_deploy: - machine: decoder + - id: rav1e-local-depth + path: dora-rav1e + build: cargo build -p dora-rav1e --release inputs: - image: rav1e-local/image - depth: rav1e-local/depth + depth: camera/depth outputs: - - image - depth - - - id: rav1e-remote - path: dora-rav1e - build: cargo build -p dora-rav1e --release - _unstable_deploy: - machine: decoder + env: + RAV1E_SPEED: 10 + - id: dav1d-local-depth + path: dora-dav1d + build: cargo build -p dora-dav1d --release inputs: - image: dav1d-remote/image - depth: dav1d-remote/depth + depth: rav1e-local-depth/depth outputs: - - image - depth - - id: dav1d-local path: dora-dav1d build: cargo build -p dora-dav1d --release - _unstable_deploy: - machine: encoder-ios inputs: - image: rav1e-remote/image - depth: rav1e-remote/depth + image: rav1e-local/image outputs: - image - - depth - id: plot build: pip install -e ../../node-hub/dora-rerun path: dora-rerun - _unstable_deploy: - machine: encoder-ios inputs: - image: dav1d-remote/image - depth: dav1d-remote/depth + image: dav1d-local/image + depth: dav1d-local-depth/depth diff --git a/node-hub/dora-dav1d/src/lib.rs b/node-hub/dora-dav1d/src/lib.rs index a829f5aa..b6aeed30 100644 --- a/node-hub/dora-dav1d/src/lib.rs +++ b/node-hub/dora-dav1d/src/lib.rs @@ -1,3 +1,5 @@ +use std::env::var; + use dav1d::Settings; use dora_node_api::{arrow::array::UInt8Array, DoraNode, Event, IntoArrow}; use eyre::{Context, Result}; @@ -49,6 +51,8 @@ pub fn lib_main() -> Result<()> { let (mut node, mut events) = DoraNode::init_from_env().context("Could not initialize dora node")?; + let output_encoding = var("ENCODING").unwrap_or("bgr8".to_string()); + loop { match events.recv() { Some(Event::Input { @@ -82,23 +86,96 @@ pub fn lib_main() -> Result<()> { let y = p.plane(dav1d::PlanarImageComponent::Y); let u = p.plane(dav1d::PlanarImageComponent::U); let v = p.plane(dav1d::PlanarImageComponent::V); - let y = yuv420_to_bgr(&y, &u, &v, p.width(), p.height()); - let arrow = y.into_arrow(); - metadata.parameters.insert( - "encoding".to_string(), - dora_node_api::Parameter::String("bgr8".to_string()), - ); - node.send_output(id, metadata.parameters, arrow).unwrap(); + match output_encoding.as_str() { + "yuv420" => { + let mut y = y.to_vec(); + let mut u = u.to_vec(); + let mut v = v.to_vec(); + y.append(&mut u); + y.append(&mut v); + let arrow = y.into_arrow(); + metadata.parameters.insert( + "encoding".to_string(), + dora_node_api::Parameter::String( + "yuv420".to_string(), + ), + ); + metadata.parameters.insert( + "width".to_string(), + dora_node_api::Parameter::Integer( + p.width() as i64 + ), + ); + metadata.parameters.insert( + "height".to_string(), + dora_node_api::Parameter::Integer( + p.height() as i64 + ), + ); + + node.send_output(id, metadata.parameters, arrow) + .unwrap(); + } + "bgr8" => { + let y = yuv420_to_bgr( + &y, + &u, + &v, + p.width(), + p.height(), + ); + let arrow = y.into_arrow(); + metadata.parameters.insert( + "encoding".to_string(), + dora_node_api::Parameter::String( + "bgr8".to_string(), + ), + ); + node.send_output(id, metadata.parameters, arrow) + .unwrap(); + } + _ => { + warn!( + "Unsupported output encoding {}", + output_encoding + ); + continue; + } + } } dav1d::PixelLayout::I400 => { let y = p.plane(dav1d::PlanarImageComponent::Y); - let vec16: Vec = bytemuck::cast_slice(&y).to_vec(); - let arrow = vec16.into_arrow(); - metadata.parameters.insert( - "encoding".to_string(), - dora_node_api::Parameter::String("mono16".to_string()), - ); - node.send_output(id, metadata.parameters, arrow).unwrap(); + match p.bit_depth() { + 8 => { + let y = y.to_vec(); + let arrow = y.into_arrow(); + metadata.parameters.insert( + "encoding".to_string(), + dora_node_api::Parameter::String( + "mono8".to_string(), + ), + ); + node.send_output(id, metadata.parameters, arrow) + .unwrap(); + } + 10 | 12 => { + let vec16: Vec = + bytemuck::cast_slice(&y).to_vec(); + let arrow = vec16.into_arrow(); + metadata.parameters.insert( + "encoding".to_string(), + dora_node_api::Parameter::String( + "mono16".to_string(), + ), + ); + node.send_output(id, metadata.parameters, arrow) + .unwrap(); + } + _ => { + warn!("Unsupported bit depth {}", p.bit_depth()); + continue; + } + } } _ => { warn!("Unsupported pixel layout"); diff --git a/node-hub/dora-ios-lidar/dora_ios_lidar/main.py b/node-hub/dora-ios-lidar/dora_ios_lidar/main.py index 8c6a0b66..13c67524 100644 --- a/node-hub/dora-ios-lidar/dora_ios_lidar/main.py +++ b/node-hub/dora-ios-lidar/dora_ios_lidar/main.py @@ -12,6 +12,7 @@ from scipy.spatial.transform import Rotation image_width = os.getenv("IMAGE_WIDTH") image_height = os.getenv("IMAGE_HEIGHT") +ROTATE = os.getenv("ROTATE") class DemoApp: @@ -100,8 +101,15 @@ class DemoApp: f_1 = intrinsic_mat[1, 1] * (int(image_width) / rgb.shape[1]) r_0 = intrinsic_mat[0, 2] * (int(image_height) / rgb.shape[0]) r_1 = intrinsic_mat[1, 2] * (int(image_width) / rgb.shape[1]) + if ROTATE == "ROTATE_90_CLOCKWISE": + rgb = cv2.rotate(rgb, cv2.ROTATE_90_CLOCKWISE) + depth = cv2.rotate(depth, cv2.ROTATE_90_CLOCKWISE) rgb = cv2.resize(rgb, (int(image_width), int(image_height))) - depth = cv2.resize(depth, (int(image_width), int(image_height))) + depth = cv2.resize( + depth, + (int(image_width), int(image_height)), + interpolation=cv2.INTER_NEAREST, + ) else: f_0 = intrinsic_mat[0, 0] f_1 = intrinsic_mat[1, 1] diff --git a/node-hub/dora-rav1e/src/lib.rs b/node-hub/dora-rav1e/src/lib.rs index 6884f7b6..2491d42f 100644 --- a/node-hub/dora-rav1e/src/lib.rs +++ b/node-hub/dora-rav1e/src/lib.rs @@ -201,6 +201,13 @@ fn send_yuv( metadata .parameters .insert("encoding".to_string(), Parameter::String("av1".to_string())); + metadata + .parameters + .insert("height".to_string(), Parameter::Integer(enc.height as i64)); + metadata + .parameters + .insert("width".to_string(), Parameter::Integer(enc.width as i64)); + let data = pkt.data; let arrow = data.into_arrow(); node.send_output(id, metadata.parameters.clone(), arrow) @@ -274,6 +281,12 @@ pub fn lib_main() -> Result<()> { height, speed_settings: SpeedSettings::from_preset(speed), low_latency: true, + chroma_sampling: color::ChromaSampling::Cs420, + color_description: Some(ColorDescription { + matrix_coefficients: MatrixCoefficients::BT709, + transfer_characteristics: color::TransferCharacteristics::BT709, + color_primaries: color::ColorPrimaries::BT709, + }), ..Default::default() }; match encoding { @@ -327,9 +340,9 @@ pub fn lib_main() -> Result<()> { 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); + let cfg = Config::new().with_encoder_config(enc.clone()); let mut ctx: Context = cfg.new_context().unwrap(); let mut f = ctx.new_frame(); diff --git a/node-hub/opencv-plot/opencv_plot/main.py b/node-hub/opencv-plot/opencv_plot/main.py index a7007d42..78b26756 100644 --- a/node-hub/opencv-plot/opencv_plot/main.py +++ b/node-hub/opencv-plot/opencv_plot/main.py @@ -1,12 +1,19 @@ """TODO: Add docstring.""" import argparse +import io import os import cv2 import numpy as np import pyarrow as pa from dora import Node +from PIL import ( + Image, +) + +if True: + import pillow_avif # noqa # noqa RUNNER_CI = True if os.getenv("CI") == "true" else False @@ -76,6 +83,7 @@ def plot_frame(plot): def yuv420p_to_bgr_opencv(yuv_array, width, height): """TODO: Add docstring.""" + yuv_array = yuv_array[: width * height * 3 // 2] yuv = yuv_array.reshape((height * 3 // 2, width)) return cv2.cvtColor(yuv, cv2.COLOR_YUV420p2RGB) @@ -145,7 +153,6 @@ def main(): encoding = metadata["encoding"] width = metadata["width"] height = metadata["height"] - if encoding == "bgr8": channels = 3 storage_type = np.uint8 @@ -181,6 +188,14 @@ def main(): img_bgr_restored = yuv420p_to_bgr_opencv(storage, width, height) plot.frame = img_bgr_restored + elif encoding == "avif": + # Convert AVIF to RGB + array = storage.to_numpy() + bytes = array.tobytes() + img = Image.open(io.BytesIO(bytes)) + img = img.convert("RGB") + plot.frame = np.array(img) + plot.frame = cv2.cvtColor(plot.frame, cv2.COLOR_RGB2BGR) else: raise RuntimeError(f"Unsupported image encoding: {encoding}") @@ -188,6 +203,7 @@ def main(): if not RUNNER_CI: if cv2.waitKey(1) & 0xFF == ord("q"): break + elif event_id == "bbox": arrow_bbox = event["value"][0] bbox_format = event["metadata"]["format"] diff --git a/node-hub/opencv-plot/pyproject.toml b/node-hub/opencv-plot/pyproject.toml index 95f253ee..6b4412a4 100644 --- a/node-hub/opencv-plot/pyproject.toml +++ b/node-hub/opencv-plot/pyproject.toml @@ -10,7 +10,13 @@ description = "Dora Node for plotting text and bbox on image with OpenCV" requires-python = ">=3.8" -dependencies = ["dora-rs >= 0.3.9", "numpy < 2.0.0", "opencv-python >= 4.1.1"] +dependencies = [ + "dora-rs >= 0.3.9", + "numpy < 2.0.0", + "opencv-python >= 4.1.1", + "pillow-avif-plugin>=1.5.1", + "pillow>=10.4.0", +] [dependency-groups] dev = ["pytest >=8.1.1", "ruff >=0.9.1"] From 7389fa419593be3877c86db33b6ffdae6fe26587 Mon Sep 17 00:00:00 2001 From: haixuanTao Date: Wed, 16 Apr 2025 14:02:10 +0200 Subject: [PATCH 3/4] Improve depth irregularities when compressing --- node-hub/dora-rav1e/src/lib.rs | 216 +++++++++++++++++---------------- 1 file changed, 110 insertions(+), 106 deletions(-) diff --git a/node-hub/dora-rav1e/src/lib.rs b/node-hub/dora-rav1e/src/lib.rs index 2491d42f..7ea00780 100644 --- a/node-hub/dora-rav1e/src/lib.rs +++ b/node-hub/dora-rav1e/src/lib.rs @@ -9,7 +9,8 @@ use std::env::var; -use dora_node_api::arrow::array::{UInt16Array, UInt8Array}; +use dora_node_api::arrow::array::AsArray; +use dora_node_api::arrow::datatypes::{UInt16Type, UInt8Type}; use dora_node_api::dora_core::config::DataId; use dora_node_api::{DoraNode, Event, IntoArrow, Metadata, Parameter}; use eyre::{Context as EyreContext, Result}; @@ -29,10 +30,12 @@ pub fn fill_zeros_toward_center_y_plane_in_place(y: &mut [u16], width: usize, he // --- Fill left half (left to center) let mut last = 0u16; - for col in 0..center { + for col in (0..center).rev() { let idx = row_start + col; if y[idx] == 0 { y[idx] = last; + } else if y[idx] > 4096 { + y[idx] = 4096; } else { last = y[idx]; } @@ -40,10 +43,12 @@ pub fn fill_zeros_toward_center_y_plane_in_place(y: &mut [u16], width: usize, he // --- Fill right half (right to center) last = 0u16; - for col in (center..width).rev() { + for col in center..width { let idx = row_start + col; if y[idx] == 0 { y[idx] = last; + } else if y[idx] > 4096 { + y[idx] = 4096; } else { last = y[idx]; } @@ -282,11 +287,6 @@ pub fn lib_main() -> Result<()> { speed_settings: SpeedSettings::from_preset(speed), low_latency: true, chroma_sampling: color::ChromaSampling::Cs420, - color_description: Some(ColorDescription { - matrix_coefficients: MatrixCoefficients::BT709, - transfer_characteristics: color::TransferCharacteristics::BT709, - color_primaries: color::ColorPrimaries::BT709, - }), ..Default::default() }; match encoding { @@ -298,117 +298,121 @@ pub fn lib_main() -> Result<()> { } 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); - send_yuv( - &y, - &u, - &v, - enc, - width, - height, - &mut node, - id, - &mut metadata, - &output_encoding, - ); + if let Some(buffer) = data.as_primitive_opt::() { + let buffer: Vec = buffer.values().to_vec(); + let (y, u, v) = bgr8_to_yuv420(buffer, width, height); + 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); - 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 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); + if let Some(buffer) = data.as_primitive_opt::() { + let buffer = buffer.values(); //.to_vec(); + + let (y, u, v) = get_yuv_planes(buffer, width, height); + send_yuv( + &y, + &u, + &v, + enc, + width, + height, + &mut node, + id, + &mut metadata, + &output_encoding, + ); } + } else if encoding == "mono16" || encoding == "z16" { + if let Some(buffer) = data.as_primitive_opt::() { + 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 bytes: &[u8] = &bytemuck::cast_slice(&buffer); + let bytes: &[u8] = &bytemuck::cast_slice(&buffer); - let cfg = Config::new().with_encoder_config(enc.clone()); - let mut ctx: Context = cfg.new_context().unwrap(); - let mut f = ctx.new_frame(); + 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; - // Multiply by 2 the stride as it is going to be width * 2 as we're converting 16-bit to 2*8-bit. - f.planes[0].copy_from_raw_u8(bytes, stride * 2, 2); + let xdec = f.planes[0].cfg.xdec; + let stride = (width + xdec) >> xdec; + // Multiply by 2 the stride as it is going to be width * 2 as we're converting 16-bit to 2*8-bit. + f.planes[0].copy_from_raw_u8(bytes, stride * 2, 2); - 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) => { - let data = pkt.data; - match output_encoding.as_str() { - "avif" => { - warn!("avif encoding not supported for mono16"); + match ctx.send_frame(f) { + Ok(_) => {} + Err(e) => match e { + EncoderStatus::EnoughData => { + warn!("Unable to send frame "); } _ => { - 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(); + warn!("Unable to send frame "); } - } + }, } - Err(e) => match e { - EncoderStatus::LimitReached => {} - EncoderStatus::Encoded => {} - EncoderStatus::NeedMoreData => {} - _ => { - panic!("Unable to receive packet",); + ctx.flush(); + match ctx.receive_packet() { + Ok(pkt) => { + let data = pkt.data; + 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 => {} + EncoderStatus::Encoded => {} + EncoderStatus::NeedMoreData => {} + _ => { + panic!("Unable to receive packet",); + } + }, + } } } else if encoding == "rgb8" { - let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); - let buffer: Vec = buffer.values().to_vec(); - 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); - send_yuv( - &y, - &u, - &v, - enc, - width, - height, - &mut node, - id, - &mut metadata, - &output_encoding, - ); + if let Some(buffer) = data.as_primitive_opt::() { + let buffer: Vec = buffer.values().to_vec(); + 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); + send_yuv( + &y, + &u, + &v, + enc, + width, + height, + &mut node, + id, + &mut metadata, + &output_encoding, + ); + } } else { unimplemented!("We haven't worked on additional encodings."); } From 4bdeca5dd58c7ba7f6bb37b93710f0ab05ed4329 Mon Sep 17 00:00:00 2001 From: haixuanTao Date: Fri, 18 Apr 2025 11:12:50 +0200 Subject: [PATCH 4/4] Fix typo --- node-hub/dora-rav1e/src/lib.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/node-hub/dora-rav1e/src/lib.rs b/node-hub/dora-rav1e/src/lib.rs index 7ea00780..22e43180 100644 --- a/node-hub/dora-rav1e/src/lib.rs +++ b/node-hub/dora-rav1e/src/lib.rs @@ -117,7 +117,7 @@ fn send_yuv( node: &mut DoraNode, id: DataId, metadata: &mut Metadata, - output_enconding: &str, + output_encoding: &str, ) -> () { // Create a new Arrow array for the YUV420 data let cfg = Config::new().with_encoder_config(enc.clone()); @@ -147,7 +147,7 @@ fn send_yuv( } ctx.flush(); match ctx.receive_packet() { - Ok(pkt) => match output_enconding { + Ok(pkt) => match output_encoding { "avif" => { let data = pkt.data.clone(); metadata.parameters.insert(