From 53d6b61839a2118ff8d7ea9e23ee818e5c5f91e6 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Sun, 6 Apr 2025 19:03:29 +0200 Subject: [PATCH] Push dav1d and rav1e on pypi --- Cargo.lock | 2 + node-hub/dora-dav1d/Cargo.toml | 16 ++ node-hub/dora-dav1d/pyproject.toml | 30 +++ node-hub/dora-dav1d/src/lib.rs | 129 ++++++++++ node-hub/dora-dav1d/src/main.rs | 109 +-------- node-hub/dora-rav1e/Cargo.toml | 17 ++ node-hub/dora-rav1e/pyproject.toml | 30 +++ node-hub/dora-rav1e/src/lib.rs | 371 +++++++++++++++++++++++++++++ node-hub/dora-rav1e/src/main.rs | 351 +-------------------------- 9 files changed, 599 insertions(+), 456 deletions(-) create mode 100644 node-hub/dora-dav1d/pyproject.toml create mode 100644 node-hub/dora-dav1d/src/lib.rs create mode 100644 node-hub/dora-rav1e/pyproject.toml create mode 100644 node-hub/dora-rav1e/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index edea5544..4001c410 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3183,6 +3183,7 @@ dependencies = [ "dora-node-api", "eyre", "log", + "pyo3", "structopt", ] @@ -3441,6 +3442,7 @@ dependencies = [ "dora-node-api", "eyre", "log", + "pyo3", "rav1e", ] diff --git a/node-hub/dora-dav1d/Cargo.toml b/node-hub/dora-dav1d/Cargo.toml index af76e7e6..a0475f23 100644 --- a/node-hub/dora-dav1d/Cargo.toml +++ b/node-hub/dora-dav1d/Cargo.toml @@ -4,6 +4,10 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[features] +default = [] +python = ["pyo3"] + [dependencies] dav1d = "0.10" bitstream-io = "2.0" @@ -12,3 +16,15 @@ structopt = "0.3" dora-node-api = { workspace = true, features = ["tracing"] } eyre = "0.6.8" bytemuck = "1.7.0" +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "eyre", + "generate-import-lib", +], optional = true } + + +[lib] +name = "dora_dav1d" +path = "src/lib.rs" +crate-type = ["lib", "cdylib"] diff --git a/node-hub/dora-dav1d/pyproject.toml b/node-hub/dora-dav1d/pyproject.toml new file mode 100644 index 00000000..e11909e7 --- /dev/null +++ b/node-hub/dora-dav1d/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["maturin>=0.13.2"] +build-backend = "maturin" + +[project] +name = "dora-dav1d" +dynamic = ["version"] +license = { text = "MIT" } +requires-python = ">=3.8" + +dependencies = [ + "maturin>=1.8.2", +] + +scripts = { "dora-dav1d" = "dora_dav1d:py_main" } + +[tool.maturin] +features = ["python", "pyo3/extension-module"] + +[tool.ruff.lint] +extend-select = [ + "D", # pydocstyle + "UP", # Ruff's UP rule + "PERF", # Ruff's PERF rule + "RET", # Ruff's RET rule + "RSE", # Ruff's RSE rule + "NPY", # Ruff's NPY rule + "N", # Ruff's N rule + "I", # Ruff's I rule +] diff --git a/node-hub/dora-dav1d/src/lib.rs b/node-hub/dora-dav1d/src/lib.rs new file mode 100644 index 00000000..6f019381 --- /dev/null +++ b/node-hub/dora-dav1d/src/lib.rs @@ -0,0 +1,129 @@ +use dav1d::Settings; +use dora_node_api::{arrow::array::UInt8Array, DoraNode, Event, IntoArrow}; +use eyre::{Context, Result}; +use log::warn; + +fn yuv420_to_bgr( + y_plane: &[u8], + u_plane: &[u8], + v_plane: &[u8], + 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 { + for i in 0..width { + let y_idx = j * width + i; // Index in Y plane + let uv_idx = (j / 2) * (width / 2) + (i / 2); // Index in U/V planes + + let y = y_plane[y_idx] as f32; + let u = u_plane[uv_idx] as f32 - 128.0; + let v = v_plane[uv_idx] as f32 - 128.0; + + // Convert YUV to RGB using BT.601 standard formula + let r = (y + 1.402 * v).clamp(0.0, 255.0) as u8; + let g = (y - 0.344136 * u - 0.714136 * v).clamp(0.0, 255.0) as u8; + let b = (y + 1.772 * u).clamp(0.0, 255.0) as u8; + + // Set the RGB values in the output buffer + let rgb_idx = y_idx * 3; + rgb_data[rgb_idx] = b; + rgb_data[rgb_idx + 1] = g; + rgb_data[rgb_idx + 2] = r; + } + } + + rgb_data +} + +pub fn lib_main() -> Result<()> { + let mut settings = Settings::new(); + // settings.set_n_threads(16); + settings.set_max_frame_delay(1); + let mut dec = + dav1d::Decoder::with_settings(&settings).expect("failed to create decoder instance"); + + let (mut node, mut events) = + DoraNode::init_from_env().context("Could not initialize dora node")?; + + loop { + match events.recv() { + Some(Event::Input { + id, + data, + mut metadata, + }) => { + 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(()) +} + +#[cfg(feature = "python")] +use pyo3::{ + pyfunction, pymodule, + types::{PyModule, PyModuleMethods}, + wrap_pyfunction, Bound, PyResult, Python, +}; + +#[cfg(feature = "python")] +#[pyfunction] +fn py_main(_py: Python) -> eyre::Result<()> { + lib_main() +} + +#[cfg(feature = "python")] +#[pymodule] +fn dora_kit_car(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(py_main, &m)?)?; + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + Ok(()) +} diff --git a/node-hub/dora-dav1d/src/main.rs b/node-hub/dora-dav1d/src/main.rs index a31b48ed..ecebf969 100644 --- a/node-hub/dora-dav1d/src/main.rs +++ b/node-hub/dora-dav1d/src/main.rs @@ -1,108 +1,3 @@ -use dav1d::Settings; -use dora_node_api::{arrow::array::UInt8Array, DoraNode, Event, IntoArrow}; -use eyre::{Context, Result}; -use log::warn; - -fn yuv420_to_bgr( - y_plane: &[u8], - u_plane: &[u8], - v_plane: &[u8], - 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 { - for i in 0..width { - let y_idx = j * width + i; // Index in Y plane - let uv_idx = (j / 2) * (width / 2) + (i / 2); // Index in U/V planes - - let y = y_plane[y_idx] as f32; - let u = u_plane[uv_idx] as f32 - 128.0; - let v = v_plane[uv_idx] as f32 - 128.0; - - // Convert YUV to RGB using BT.601 standard formula - let r = (y + 1.402 * v).clamp(0.0, 255.0) as u8; - let g = (y - 0.344136 * u - 0.714136 * v).clamp(0.0, 255.0) as u8; - let b = (y + 1.772 * u).clamp(0.0, 255.0) as u8; - - // Set the RGB values in the output buffer - let rgb_idx = y_idx * 3; - rgb_data[rgb_idx] = b; - rgb_data[rgb_idx + 1] = g; - rgb_data[rgb_idx + 2] = r; - } - } - - rgb_data -} - -fn main() -> Result<()> { - let mut settings = Settings::new(); - // settings.set_n_threads(16); - settings.set_max_frame_delay(1); - let mut dec = - dav1d::Decoder::with_settings(&settings).expect("failed to create decoder instance"); - - let (mut node, mut events) = - DoraNode::init_from_env().context("Could not initialize dora node")?; - - loop { - match events.recv() { - Some(Event::Input { - id, - data, - mut metadata, - }) => { - 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(()) +fn main() -> eyre::Result<()> { + dora_dav1d::lib_main() } diff --git a/node-hub/dora-rav1e/Cargo.toml b/node-hub/dora-rav1e/Cargo.toml index cdd4249e..20f54934 100644 --- a/node-hub/dora-rav1e/Cargo.toml +++ b/node-hub/dora-rav1e/Cargo.toml @@ -9,9 +9,26 @@ repository.workspace = true # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[features] +default = [] +python = ["pyo3"] + [dependencies] rav1e = { version = "0.7.1", features = ["serialize"] } dora-node-api = { workspace = true, features = ["tracing"] } eyre = "0.6.8" log = "0.4" bytemuck = "1.20" +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "eyre", + "generate-import-lib", +], optional = true } + + +[lib] +name = "dora_rav1e" +path = "src/lib.rs" +crate-type = ["lib", "cdylib"] + diff --git a/node-hub/dora-rav1e/pyproject.toml b/node-hub/dora-rav1e/pyproject.toml new file mode 100644 index 00000000..d5aed30a --- /dev/null +++ b/node-hub/dora-rav1e/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["maturin>=0.13.2"] +build-backend = "maturin" + +[project] +name = "dora-rav1e" +dynamic = ["version"] +license = { text = "MIT" } +requires-python = ">=3.8" + +dependencies = [ + "maturin>=1.8.2", +] + +scripts = { "dora-rav1e" = "dora_rav1e:py_main" } + +[tool.maturin] +features = ["python", "pyo3/extension-module"] + +[tool.ruff.lint] +extend-select = [ + "D", # pydocstyle + "UP", # Ruff's UP rule + "PERF", # Ruff's PERF rule + "RET", # Ruff's RET rule + "RSE", # Ruff's RSE rule + "NPY", # Ruff's NPY rule + "N", # Ruff's N rule + "I", # Ruff's I rule +] diff --git a/node-hub/dora-rav1e/src/lib.rs b/node-hub/dora-rav1e/src/lib.rs new file mode 100644 index 00000000..149b39f7 --- /dev/null +++ b/node-hub/dora-rav1e/src/lib.rs @@ -0,0 +1,371 @@ +// Copyright (c) 2019-2022, The rav1e contributors. All rights reserved +// +// This source code is subject to the terms of the BSD 2 Clause License and +// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License +// was not distributed with this source code in the LICENSE file, you can +// obtain it at www.aomedia.org/license/software. If the Alliance for Open +// Media Patent License 1.0 was not distributed with this source code in the +// PATENTS file, you can obtain it at www.aomedia.org/license/patent. + +use std::env::var; + +use dora_node_api::arrow::array::{UInt16Array, UInt8Array}; +use dora_node_api::{DoraNode, Event, IntoArrow, Parameter}; +use eyre::{Context as EyreContext, Result}; +use log::warn; +// Encode the same tiny blank frame 30 times +use rav1e::config::SpeedSettings; + +use rav1e::*; + +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)]; + let mut v_plane = vec![0; (width / 2) * (height / 2)]; + + for y in 0..height { + for x in 0..width { + let bgr_index = (y * width + x) * 3; + let b = bgr_data[bgr_index] as f32; + let g = bgr_data[bgr_index + 1] as f32; + let r = bgr_data[bgr_index + 2] as f32; + + // Corrected YUV conversion formulas + let y_value = (0.299 * r + 0.587 * g + 0.114 * b).clamp(0.0, 255.0); + let u_value = (-0.14713 * r - 0.28886 * g + 0.436 * b + 128.0).clamp(0.0, 255.0); + let v_value = (0.615 * r - 0.51499 * g - 0.10001 * b + 128.0).clamp(0.0, 255.0); + + let y_index = y * width + x; + y_plane[y_index] = y_value.round() as u8; + + if x % 2 == 0 && y % 2 == 0 { + let uv_index = (y / 2) * (width / 2) + (x / 2); + u_plane[uv_index] = u_value.round() as u8; + v_plane[uv_index] = v_value.round() as u8; + } + } + } + + (y_plane, u_plane, v_plane) +} + +fn get_yuv_planes(buffer: &[u8], width: usize, height: usize) -> (&[u8], &[u8], &[u8]) { + // Calculate sizes of Y, U, and V planes for YUV420 format + let y_size = width * height; // Y has full resolution + let uv_width = width / 2; // U and V are subsampled by 2 in both dimensions + let uv_height = height / 2; // U and V are subsampled by 2 in both dimensions + let uv_size = uv_width * uv_height; // Size for U and V planes + + // Ensure the buffer has the correct size + // if buffer.len() != y_size + 2 * uv_size { + // panic!("Invalid buffer size for the given width and height!"); + // } + + // Extract Y, U, and V planes + let y_plane = &buffer[0..y_size]; //.to_vec(); + let u_plane = &buffer[y_size..y_size + uv_size]; //.to_vec(); + let v_plane = &buffer[y_size + uv_size..]; //.to_vec(); + + (y_plane, u_plane, v_plane) +} + +pub fn lib_main() -> Result<()> { + let mut height = std::env::var("IMAGE_HEIGHT") + .unwrap_or_else(|_| "480".to_string()) + .parse() + .unwrap(); + let mut width = std::env::var("IMAGE_WIDTH") + .unwrap_or_else(|_| "640".to_string()) + .parse() + .unwrap(); + let speed = var("RAV1E_SPEED").map(|s| s.parse().unwrap()).unwrap_or(10); + let mut enc = EncoderConfig { + width, + height, + speed_settings: SpeedSettings::from_preset(speed), + low_latency: true, + ..Default::default() + }; + let cfg = Config::new().with_encoder_config(enc.clone()); + cfg.validate()?; + + let (mut node, mut events) = + DoraNode::init_from_env().context("Could not initialize dora node")?; + + loop { + match events.recv() { + Some(Event::Input { + id, + data, + mut metadata, + }) => { + if let Some(Parameter::Integer(h)) = metadata.parameters.get("height") { + height = *h as usize; + }; + if let Some(Parameter::Integer(w)) = metadata.parameters.get("width") { + width = *w as usize; + }; + let encoding = if let Some(Parameter::String(encoding)) = + metadata.parameters.get("encoding") + { + encoding + } else { + "bgr8" + }; + enc = EncoderConfig { + width, + height, + speed_settings: SpeedSettings::from_preset(speed), + low_latency: true, + ..Default::default() + }; + match encoding { + "mono16" => { + enc.bit_depth = 12; + enc.chroma_sampling = color::ChromaSampling::Cs400; + } + _ => {} + } + + 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",); + } + }, + } + } 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" { + 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 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 "); + } + }, + } + 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 == "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); + + // 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",); + } + }, + } + } else { + unimplemented!("We haven't worked on additional encodings."); + } + } + Some(Event::Error(_e)) => { + continue; + } + _ => break, + }; + } + + Ok(()) +} + +#[cfg(feature = "python")] +use pyo3::{ + pyfunction, pymodule, + types::{PyModule, PyModuleMethods}, + wrap_pyfunction, Bound, PyResult, Python, +}; + +#[cfg(feature = "python")] +#[pyfunction] +fn py_main(_py: Python) -> eyre::Result<()> { + lib_main() +} + +#[cfg(feature = "python")] +#[pymodule] +fn dora_kit_car(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(py_main, &m)?)?; + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + Ok(()) +} diff --git a/node-hub/dora-rav1e/src/main.rs b/node-hub/dora-rav1e/src/main.rs index 076bee23..2e0f8be9 100644 --- a/node-hub/dora-rav1e/src/main.rs +++ b/node-hub/dora-rav1e/src/main.rs @@ -1,350 +1,3 @@ -// Copyright (c) 2019-2022, The rav1e contributors. All rights reserved -// -// This source code is subject to the terms of the BSD 2 Clause License and -// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License -// was not distributed with this source code in the LICENSE file, you can -// obtain it at www.aomedia.org/license/software. If the Alliance for Open -// Media Patent License 1.0 was not distributed with this source code in the -// PATENTS file, you can obtain it at www.aomedia.org/license/patent. - -use std::env::var; - -use dora_node_api::arrow::array::{UInt16Array, UInt8Array}; -use dora_node_api::{DoraNode, Event, IntoArrow, Parameter}; -use eyre::{Context as EyreContext, Result}; -use log::warn; -// Encode the same tiny blank frame 30 times -use rav1e::config::SpeedSettings; - -use rav1e::*; - -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)]; - let mut v_plane = vec![0; (width / 2) * (height / 2)]; - - for y in 0..height { - for x in 0..width { - let bgr_index = (y * width + x) * 3; - let b = bgr_data[bgr_index] as f32; - let g = bgr_data[bgr_index + 1] as f32; - let r = bgr_data[bgr_index + 2] as f32; - - // Corrected YUV conversion formulas - let y_value = (0.299 * r + 0.587 * g + 0.114 * b).clamp(0.0, 255.0); - let u_value = (-0.14713 * r - 0.28886 * g + 0.436 * b + 128.0).clamp(0.0, 255.0); - let v_value = (0.615 * r - 0.51499 * g - 0.10001 * b + 128.0).clamp(0.0, 255.0); - - let y_index = y * width + x; - y_plane[y_index] = y_value.round() as u8; - - if x % 2 == 0 && y % 2 == 0 { - let uv_index = (y / 2) * (width / 2) + (x / 2); - u_plane[uv_index] = u_value.round() as u8; - v_plane[uv_index] = v_value.round() as u8; - } - } - } - - (y_plane, u_plane, v_plane) -} - -fn get_yuv_planes(buffer: &[u8], width: usize, height: usize) -> (&[u8], &[u8], &[u8]) { - // Calculate sizes of Y, U, and V planes for YUV420 format - let y_size = width * height; // Y has full resolution - let uv_width = width / 2; // U and V are subsampled by 2 in both dimensions - let uv_height = height / 2; // U and V are subsampled by 2 in both dimensions - let uv_size = uv_width * uv_height; // Size for U and V planes - - // Ensure the buffer has the correct size - // if buffer.len() != y_size + 2 * uv_size { - // panic!("Invalid buffer size for the given width and height!"); - // } - - // Extract Y, U, and V planes - let y_plane = &buffer[0..y_size]; //.to_vec(); - let u_plane = &buffer[y_size..y_size + uv_size]; //.to_vec(); - let v_plane = &buffer[y_size + uv_size..]; //.to_vec(); - - (y_plane, u_plane, v_plane) -} - -fn main() -> Result<()> { - let mut height = std::env::var("IMAGE_HEIGHT") - .unwrap_or_else(|_| "480".to_string()) - .parse() - .unwrap(); - let mut width = std::env::var("IMAGE_WIDTH") - .unwrap_or_else(|_| "640".to_string()) - .parse() - .unwrap(); - let speed = var("RAV1E_SPEED").map(|s| s.parse().unwrap()).unwrap_or(10); - let mut enc = EncoderConfig { - width, - height, - speed_settings: SpeedSettings::from_preset(speed), - low_latency: true, - ..Default::default() - }; - let cfg = Config::new().with_encoder_config(enc.clone()); - cfg.validate()?; - - let (mut node, mut events) = - DoraNode::init_from_env().context("Could not initialize dora node")?; - - loop { - match events.recv() { - Some(Event::Input { - id, - data, - mut metadata, - }) => { - if let Some(Parameter::Integer(h)) = metadata.parameters.get("height") { - height = *h as usize; - }; - if let Some(Parameter::Integer(w)) = metadata.parameters.get("width") { - width = *w as usize; - }; - let encoding = if let Some(Parameter::String(encoding)) = - metadata.parameters.get("encoding") - { - encoding - } else { - "bgr8" - }; - enc = EncoderConfig { - width, - height, - speed_settings: SpeedSettings::from_preset(speed), - low_latency: true, - ..Default::default() - }; - match encoding { - "mono16" => { - enc.bit_depth = 12; - enc.chroma_sampling = color::ChromaSampling::Cs400; - } - _ => {} - } - - 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",); - } - }, - } - } 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" { - 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 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 "); - } - }, - } - 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 == "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); - - // 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",); - } - }, - } - } else { - unimplemented!("We haven't worked on additional encodings."); - } - } - Some(Event::Error(_e)) => { - continue; - } - _ => break, - }; - } - - Ok(()) +fn main() -> eyre::Result<()> { + dora_rav1e::lib_main() }