dora-rs
/
dora
Not watched
1
0
Fork 0
Code
Releases 124 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
Browse Source

Fix mono16 encoding

tags/v0.3.11-rc1
haixuanTao haixuantao 9 months ago
parent
73b9823dba
commit
61fb91eb53
5 changed files with 165 additions and 122 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +2
    -0
      Cargo.lock
  2. +45
    -55
      node-hub/dora-dav1d/src/main.rs
  3. +1
    -0
      node-hub/dora-rav1e/Cargo.toml
  4. +10
    -16
      node-hub/dora-rav1e/src/main.rs
  5. +107
    -51
      node-hub/dora-rerun/src/lib.rs

+ 2
- 0
Cargo.lock View File

@@ -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",


+ 45
- 55
node-hub/dora-dav1d/src/main.rs View File

@@ -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<u8> {
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::<UInt8Array>().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<u16> = 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::<UInt8Array>() {
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<u16> = 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(())
}

+ 1
- 0
node-hub/dora-rav1e/Cargo.toml View File

@@ -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"

+ 10
- 16
node-hub/dora-rav1e/src/main.rs View File

@@ -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<u16> = 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<u8> = 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.");
}


+ 107
- 51
node-hub/dora-rerun/src/lib.rs View File

@@ -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::<Vec<_>>()
} else {
color_buffer
.chunks(3)
.map(|x| rerun::Color::from_rgb(x[0], x[1], x[2]))
.collect::<Vec<_>>()
};
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::<Vec<_>>()
} else {
color_buffer
.chunks(3)
.map(|x| rerun::Color::from_rgb(x[0], x[1], x[2]))
.collect::<Vec<_>>()
};
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::<Vec<_>>()
} else {
color_buffer
.chunks(3)
.map(|x| rerun::Color::from_rgb(x[0], x[1], x[2]))
.collect::<Vec<_>>()
};
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::<Float32Type>() {
let data = data
.iter()


Write Preview
Loading…
Cancel
Save