diff --git a/Cargo.lock b/Cargo.lock index 40d2d356..dbb2f973 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3178,6 +3178,7 @@ name = "dora-dav1d" version = "0.0.0" dependencies = [ "bitstream-io", + "bytemuck", "dav1d", "dora-node-api", "eyre", @@ -3436,6 +3437,7 @@ dependencies = [ name = "dora-rav1e" version = "0.3.10" dependencies = [ + "bytemuck", "dora-node-api", "eyre", "log", diff --git a/node-hub/dora-dav1d/src/main.rs b/node-hub/dora-dav1d/src/main.rs index 44f0f55c..a31b48ed 100644 --- a/node-hub/dora-dav1d/src/main.rs +++ b/node-hub/dora-dav1d/src/main.rs @@ -1,5 +1,5 @@ use dav1d::Settings; -use dora_node_api::{arrow::array::UInt8Array, DoraNode, Event, IntoArrow, Parameter}; +use dora_node_api::{arrow::array::UInt8Array, DoraNode, Event, IntoArrow}; use eyre::{Context, Result}; use log::warn; @@ -7,9 +7,11 @@ fn yuv420_to_bgr( y_plane: &[u8], u_plane: &[u8], v_plane: &[u8], - width: usize, - height: usize, + width: u32, + height: u32, ) -> Vec { + let width = width as usize; + let height = height as usize; let mut rgb_data = vec![0u8; width * height * 3]; // Output RGB data buffer for j in 0..height { @@ -54,65 +56,53 @@ fn main() -> Result<()> { data, mut metadata, }) => { - let data = data.as_any().downcast_ref::().unwrap(); - let data = data.values().clone(); - - match dec.send_data(data, None, None, None) { - Err(e) => { - warn!("Error sending data to the decoder: {}", e); - } - Ok(()) => { - if let Ok(p) = dec.get_picture() { - let height = if let Some(Parameter::Integer(h)) = - metadata.parameters.get("height") - { - *h as usize - } else { - 640 - }; - let width = if let Some(Parameter::Integer(w)) = - metadata.parameters.get("width") - { - *w as usize - } else { - 480 - }; - match p.pixel_layout() { - dav1d::PixelLayout::I420 => { - 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, width, 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(); - } - 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(); - } - _ => { - warn!("Unsupported pixel layout"); - continue; - } - }; + if let Some(data) = data.as_any().downcast_ref::() { + let data = data.values().clone(); + match dec.send_data(data, None, None, None) { + Err(e) => { + warn!("Error sending data to the decoder: {}", e); + } + Ok(()) => { + if let Ok(p) = dec.get_picture() { + match p.pixel_layout() { + dav1d::PixelLayout::I420 => { + 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(); + } + 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(); + } + _ => { + warn!("Unsupported pixel layout"); + continue; + } + }; + } } } + } else { + warn!("Unsupported data type {}", data.data_type()); + continue; } } None => break, Some(_) => break, } } - Ok(()) } diff --git a/node-hub/dora-rav1e/Cargo.toml b/node-hub/dora-rav1e/Cargo.toml index 10a6bcda..cdd4249e 100644 --- a/node-hub/dora-rav1e/Cargo.toml +++ b/node-hub/dora-rav1e/Cargo.toml @@ -14,3 +14,4 @@ rav1e = { version = "0.7.1", features = ["serialize"] } dora-node-api = { workspace = true, features = ["tracing"] } eyre = "0.6.8" log = "0.4" +bytemuck = "1.20" diff --git a/node-hub/dora-rav1e/src/main.rs b/node-hub/dora-rav1e/src/main.rs index 38e114b5..91e37973 100644 --- a/node-hub/dora-rav1e/src/main.rs +++ b/node-hub/dora-rav1e/src/main.rs @@ -93,7 +93,7 @@ fn main() -> Result<()> { DoraNode::init_from_env().context("Could not initialize dora node")?; loop { - let _buffer = match events.recv() { + match events.recv() { Some(Event::Input { id, data, @@ -154,9 +154,11 @@ fn main() -> Result<()> { Err(e) => match e { EncoderStatus::EnoughData => { warn!("Unable to send frame "); + continue; } _ => { warn!("Unable to send frame "); + continue; } }, } @@ -181,7 +183,6 @@ fn main() -> Result<()> { } }, } - vec![] } else if encoding == "yuv420" { let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); let buffer = buffer.values(); //.to_vec(); @@ -232,15 +233,19 @@ fn main() -> Result<()> { } }, } - vec![] } else if encoding == "mono16" { let buffer: &UInt16Array = data.as_any().downcast_ref().unwrap(); let buffer: &[u16] = buffer.values(); + // let buffer = shift_u16_slice_to_upper_12_bits(buffer); + let bytes: &[u8] = &bytemuck::cast_slice(&buffer); + let mut ctx: Context = cfg.new_context().unwrap(); let mut f = ctx.new_frame(); - let origin = f.planes[0].data_origin_mut(); - origin.copy_from_slice(buffer); + 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(_) => {} @@ -274,19 +279,8 @@ fn main() -> Result<()> { } }, } - vec![] } else if encoding == "rgb8" { unimplemented!("We haven't worked on additional encodings."); - let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); - let buffer: &[u8] = buffer.values(); - let mut ctx: Context = cfg.new_context().unwrap(); - let mut f = ctx.new_frame(); - - for p in &mut f.planes { - let stride = (enc.width + p.cfg.xdec) >> p.cfg.xdec; - p.copy_from_raw_u8(&buffer, stride, 1); - } - buffer.to_vec() } 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 c79d39c8..fa0b9520 100644 --- a/node-hub/dora-rerun/src/lib.rs +++ b/node-hub/dora-rerun/src/lib.rs @@ -3,10 +3,8 @@ use std::{collections::HashMap, env::VarError, path::Path}; use dora_node_api::{ - arrow::{ - array::{Array, AsArray, Float32Array, Float64Array, StringArray, UInt8Array}, - datatypes::Float32Type, - }, + arrow::array::{Array, Float32Array, Float64Array, StringArray, UInt16Array, UInt8Array}, + arrow::{array::AsArray, datatypes::Float32Type}, dora_core::config::DataId, DoraNode, Event, Parameter, }; @@ -181,54 +179,112 @@ pub fn lib_main() -> Result<()> { } else { vec![640, 480] }; - let buffer: &Float64Array = data.as_any().downcast_ref().unwrap(); - let points_3d = buffer.iter().enumerate().map(|(i, z)| { - let u = i as f32 % *width as f32; // Calculate x-coordinate (u) - let v = i as f32 / *width as f32; // Calculate y-coordinate (v) - let z = z.unwrap_or_default() as f32; + match data.data_type() { + dora_node_api::arrow::datatypes::DataType::Float64 => { + let buffer: &Float64Array = data.as_any().downcast_ref().unwrap(); - ( - (u - resolution[0] as f32) * z / focal_length[0] as f32, - (v - resolution[1] as f32) * z / focal_length[1] as f32, - z, - ) - }); - let points_3d = Points3D::new(points_3d); - if let Some(color_buffer) = image_cache.get(&id.replace("depth", "image")) { - let colors = if let Some(mask) = mask_cache.get(&id.replace("depth", "mask")) { - let mask_length = color_buffer.len() / 3; - let number_masks = mask.len() / mask_length; - color_buffer - .chunks(3) - .enumerate() - .map(|(e, x)| { - for i in 0..number_masks { - if mask[i * mask_length + e] && (e % 3 == 0) { - if i == 0 { - return rerun::Color::from_rgb(255, x[1], x[2]); - } else if i == 1 { - return rerun::Color::from_rgb(x[0], 255, x[2]); - } else if i == 2 { - return rerun::Color::from_rgb(x[0], x[1], 255); - } else { - return rerun::Color::from_rgb(x[0], 255, x[2]); + let points_3d = buffer.iter().enumerate().map(|(i, z)| { + let u = i as f32 % *width as f32; // Calculate x-coordinate (u) + let v = i as f32 / *width as f32; // Calculate y-coordinate (v) + let z = z.unwrap_or_default() as f32; + + ( + (u - resolution[0] as f32) * z / focal_length[0] as f32, + (v - resolution[1] as f32) * z / focal_length[1] as f32, + z, + ) + }); + let points_3d = Points3D::new(points_3d); + if let Some(color_buffer) = image_cache.get(&id.replace("depth", "image")) { + let colors = if let Some(mask) = + mask_cache.get(&id.replace("depth", "masks")) + { + let mask_length = color_buffer.len() / 3; + let number_masks = mask.len() / mask_length; + color_buffer + .chunks(3) + .enumerate() + .map(|(e, x)| { + for i in 0..number_masks { + if mask[i * mask_length + e] && (e % 3 == 0) { + if i == 0 { + return rerun::Color::from_rgb(255, x[1], x[2]); + } else if i == 1 { + return rerun::Color::from_rgb(x[0], 255, x[2]); + } else if i == 2 { + return rerun::Color::from_rgb(x[0], x[1], 255); + } else { + return rerun::Color::from_rgb(x[0], 255, x[2]); + } + } } - } - } - rerun::Color::from_rgb(x[0], x[1], x[2]) - }) - .collect::>() - } else { - color_buffer - .chunks(3) - .map(|x| rerun::Color::from_rgb(x[0], x[1], x[2])) - .collect::>() - }; - rec.log(id.as_str(), &points_3d.with_colors(colors)) - .context("could not log points")?; - } else { - rec.log(id.as_str(), &points_3d) - .context("could not log points")?; + rerun::Color::from_rgb(x[0], x[1], x[2]) + }) + .collect::>() + } else { + color_buffer + .chunks(3) + .map(|x| rerun::Color::from_rgb(x[0], x[1], x[2])) + .collect::>() + }; + rec.log(id.as_str(), &points_3d.with_colors(colors)) + .context("could not log points")?; + } + } + dora_node_api::arrow::datatypes::DataType::UInt16 => { + let buffer: &UInt16Array = data.as_any().downcast_ref().unwrap(); + + let points_3d = buffer.iter().enumerate().map(|(i, z)| { + let u = i as f32 % *width as f32; // Calculate x-coordinate (u) + let v = i as f32 / *width as f32; // Calculate y-coordinate (v) + let z = z.unwrap_or_default() as f32 / 1_000.; + + ( + (u - resolution[0] as f32) * z / focal_length[0] as f32, + (v - resolution[1] as f32) * z / focal_length[1] as f32, + z, + ) + }); + let points_3d = Points3D::new(points_3d); + if let Some(color_buffer) = image_cache.get(&id.replace("depth", "image")) { + let colors = if let Some(mask) = + mask_cache.get(&id.replace("depth", "masks")) + { + let mask_length = color_buffer.len() / 3; + let number_masks = mask.len() / mask_length; + color_buffer + .chunks(3) + .enumerate() + .map(|(e, x)| { + for i in 0..number_masks { + if mask[i * mask_length + e] && (e % 3 == 0) { + if i == 0 { + return rerun::Color::from_rgb(255, x[1], x[2]); + } else if i == 1 { + return rerun::Color::from_rgb(x[0], 255, x[2]); + } else if i == 2 { + return rerun::Color::from_rgb(x[0], x[1], 255); + } else { + return rerun::Color::from_rgb(x[0], 255, x[2]); + } + } + } + rerun::Color::from_rgb(x[0], x[1], x[2]) + }) + .collect::>() + } else { + color_buffer + .chunks(3) + .map(|x| rerun::Color::from_rgb(x[0], x[1], x[2])) + .collect::>() + }; + rec.log(id.as_str(), &points_3d.with_colors(colors)) + .context("could not log points")?; + } + } + _ => { + return Err(eyre!("Unsupported depth data type {}", data.data_type())); + } } } else if id.as_str().contains("text") { let buffer: StringArray = data.to_data().into(); @@ -242,7 +298,7 @@ pub fn lib_main() -> Result<()> { })?; } else if id.as_str().contains("boxes2d") { boxes2d::update_boxes2d(&rec, id, data, metadata).context("update boxes 2d")?; - } else if id.as_str().contains("mask") { + } else if id.as_str().contains("masks") { let masks = if let Some(data) = data.as_primitive_opt::() { let data = data .iter()