From a82c2a451b1c429a7fd283cdfe89be8087f5e8c5 Mon Sep 17 00:00:00 2001 From: haixuanTao Date: Wed, 26 Mar 2025 21:34:21 +0100 Subject: [PATCH] Initial support for dav1d and rav1e --- Cargo.lock | 275 +++++++++++++++++++++++++- Cargo.toml | 2 + node-hub/dora-av1-encoder/Cargo.toml | 15 ++ node-hub/dora-av1-encoder/example.py | 22 +++ node-hub/dora-av1-encoder/src/main.rs | 217 ++++++++++++++++++++ node-hub/dora-dav1d/Cargo.toml | 13 ++ node-hub/dora-dav1d/dataflow.yml | 53 +++++ node-hub/dora-dav1d/src/main.rs | 189 ++++++++++++++++++ 8 files changed, 784 insertions(+), 2 deletions(-) create mode 100644 node-hub/dora-av1-encoder/Cargo.toml create mode 100644 node-hub/dora-av1-encoder/example.py create mode 100644 node-hub/dora-av1-encoder/src/main.rs create mode 100644 node-hub/dora-dav1d/Cargo.toml create mode 100644 node-hub/dora-dav1d/dataflow.yml create mode 100644 node-hub/dora-dav1d/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 957cd526..9482a998 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -273,6 +273,15 @@ dependencies = [ "rgb", ] +[[package]] +name = "ansi_term" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d52a9bb7ec0cf484c551830a7ce27bd20d67eac647e1befb56b0be4ee39a55d2" +dependencies = [ + "winapi 0.3.9", +] + [[package]] name = "anstream" version = "0.6.18" @@ -1297,6 +1306,36 @@ dependencies = [ "num-traits", ] +[[package]] +name = "av-metrics" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "996ce95bbdb0203e5b91d4a0c9b81c0d67d11c80f884482a0c1ea19e732e3530" +dependencies = [ + "crossbeam", + "itertools 0.10.5", + "lab", + "num-traits", + "rayon", + "thiserror 1.0.66", + "v_frame", +] + +[[package]] +name = "av1-grain" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6678909d8c5d46a42abcf571271e15fdbc0a225e3646cf23762cd415046c78bf" +dependencies = [ + "anyhow", + "arrayvec", + "log", + "nom", + "num-rational", + "serde", + "v_frame", +] + [[package]] name = "axum" version = "0.7.9" @@ -1605,6 +1644,15 @@ dependencies = [ "sysinfo 0.33.1", ] +[[package]] +name = "built" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6a6c0b39c38fd754ac338b00a88066436389c0f029da5d37d1e01091d9b7c17" +dependencies = [ + "git2", +] + [[package]] name = "bumpalo" version = "3.17.0" @@ -1905,6 +1953,16 @@ dependencies = [ "uuid 1.16.0", ] +[[package]] +name = "cfg-expr" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02" +dependencies = [ + "smallvec", + "target-lexicon", +] + [[package]] name = "cfg-if" version = "0.1.10" @@ -1986,6 +2044,21 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "clap" +version = "2.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0610544180c38b88101fecf2dd634b174a62eef6946f84dfc6a7127512b381c" +dependencies = [ + "ansi_term", + "atty", + "bitflags 1.3.2", + "strsim 0.8.0", + "textwrap 0.11.0", + "unicode-width", + "vec_map", +] + [[package]] name = "clap" version = "3.2.25" @@ -2000,7 +2073,7 @@ dependencies = [ "once_cell", "strsim 0.10.0", "termcolor", - "textwrap", + "textwrap 0.16.1", ] [[package]] @@ -2026,6 +2099,15 @@ dependencies = [ "terminal_size", ] +[[package]] +name = "clap_complete" +version = "4.5.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06f5378ea264ad4f82bbc826628b5aad714a75abf6ece087e923010eb937fb6" +dependencies = [ + "clap 4.5.20", +] + [[package]] name = "clap_derive" version = "3.2.25" @@ -2937,6 +3019,15 @@ dependencies = [ "num", ] +[[package]] +name = "dora-av1-encoder" +version = "0.3.10" +dependencies = [ + "dora-node-api", + "eyre", + "rav1e", +] + [[package]] name = "dora-cli" version = "0.3.10" @@ -3049,6 +3140,18 @@ dependencies = [ "zenoh 1.3.0", ] +[[package]] +name = "dora-dav1d" +version = "0.0.0" +dependencies = [ + "bitstream-io", + "dav1d", + "dora-node-api", + "eyre", + "log", + "structopt", +] + [[package]] name = "dora-download" version = "0.3.10" @@ -4086,6 +4189,15 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "fern" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9f0c14694cbd524c8720dd69b0e3179344f04ebb5f90f2e4a440c6ea3b2f1ee" +dependencies = [ + "log", +] + [[package]] name = "ffmpeg-sidecar" version = "2.0.5" @@ -6044,6 +6156,12 @@ dependencies = [ "log", ] +[[package]] +name = "lab" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf36173d4167ed999940f804952e6b08197cae5ad5d572eb4db150ce8ad5d58f" + [[package]] name = "lazy_static" version = "1.5.0" @@ -6135,6 +6253,16 @@ version = "0.2.171" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" +[[package]] +name = "libfuzzer-sys" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf78f52d400cf2d84a3a973a78a592b4adc535739e0a5597a0da6f0c357adc75" +dependencies = [ + "arbitrary", + "cc", +] + [[package]] name = "libgit2-sys" version = "0.16.2+1.7.2" @@ -6407,6 +6535,16 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "maybe-rayon" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ea1f30cedd69f0a2954655f7188c6a834246d2bcf1e315e2ac40c4b24dc9519" +dependencies = [ + "cfg-if 1.0.0", + "rayon", +] + [[package]] name = "md5" version = "0.7.0" @@ -6918,6 +7056,15 @@ dependencies = [ "getrandom 0.2.15", ] +[[package]] +name = "nasm-rs" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe4d98d0065f4b1daf164b3eafb11974c94662e5e2396cf03f32d0bb5c17da51" +dependencies = [ + "rayon", +] + [[package]] name = "nasm-rs" version = "0.3.0" @@ -7042,6 +7189,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c96aba5aa877601bb3f6dd6a63a969e1f82e60646e81e71b14496995e9853c91" +[[package]] +name = "new_debug_unreachable" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" + [[package]] name = "newline-converter" version = "0.2.2" @@ -9092,6 +9245,55 @@ dependencies = [ "rand_core 0.9.3", ] +[[package]] +name = "rav1e" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd87ce80a7665b1cce111f8a16c1f3929f6547ce91ade6addf4ec86a8dda5ce9" +dependencies = [ + "arbitrary", + "arg_enum_proc_macro", + "arrayvec", + "av-metrics", + "av1-grain", + "bitstream-io", + "built", + "cc", + "cfg-if 1.0.0", + "clap 4.5.20", + "clap_complete", + "console", + "fern", + "interpolate_name", + "itertools 0.12.1", + "ivf", + "libc", + "libfuzzer-sys", + "log", + "maybe-rayon", + "nasm-rs 0.2.5", + "new_debug_unreachable", + "nom", + "noop_proc_macro", + "num-derive", + "num-traits", + "once_cell", + "paste", + "profiling", + "rand", + "rand_chacha", + "scan_fmt", + "serde", + "serde-big-array", + "signal-hook", + "simd_helpers", + "system-deps", + "thiserror 1.0.66", + "toml", + "v_frame", + "y4m", +] + [[package]] name = "raw-cpuid" version = "10.7.0" @@ -9829,7 +10031,7 @@ dependencies = [ "cc", "cfg-if 1.0.0", "libc", - "nasm-rs", + "nasm-rs 0.3.0", "parking_lot", "paste", "raw-cpuid 11.5.0", @@ -11591,6 +11793,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scan_fmt" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b53b0a5db882a8e2fdaae0a43f7b39e7e9082389e978398bdf223a55b581248" + [[package]] name = "schannel" version = "0.1.27" @@ -12183,6 +12391,15 @@ version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" +[[package]] +name = "simd_helpers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95890f873bec569a0362c235787f3aca6e1e887302ba4840839bcc6459c42da6" +dependencies = [ + "quote", +] + [[package]] name = "simdutf8" version = "0.1.5" @@ -12548,6 +12765,12 @@ dependencies = [ "float-cmp", ] +[[package]] +name = "strsim" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a" + [[package]] name = "strsim" version = "0.10.0" @@ -12560,6 +12783,30 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "structopt" +version = "0.3.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6b5c64445ba8094a6ab0c3cd2ad323e07171012d9c98b0b15651daf1787a10" +dependencies = [ + "clap 2.34.0", + "lazy_static", + "structopt-derive", +] + +[[package]] +name = "structopt-derive" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcb5ae327f9cc13b68763b5749770cb9e048a99bd9dfdfa58d0cf05d5f64afe0" +dependencies = [ + "heck 0.3.3", + "proc-macro-error", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "strum" version = "0.26.3" @@ -13897,6 +14144,18 @@ dependencies = [ "syn 2.0.100", ] +[[package]] +name = "v_frame" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f32aaa24bacd11e488aa9ba66369c7cd514885742c9fe08cfe85884db3e92b" +dependencies = [ + "aligned-vec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "validated_struct" version = "2.1.1" @@ -13968,6 +14227,12 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1bddf1187be692e79c5ffeab891132dfb0f236ed36a43c7ed39f1165ee20191" +[[package]] +name = "version-compare" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "852e951cb7832cb45cb1169900d19760cfa39b82bc0ea9c0e5a14ae88411c98b" + [[package]] name = "version_check" version = "0.9.5" @@ -15404,6 +15669,12 @@ version = "0.8.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +[[package]] +name = "y4m" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a5a4b21e1a62b67a2970e6831bc091d7b87e119e7f9791aef9702e3bef04448" + [[package]] name = "yaml-rust" version = "0.4.5" diff --git a/Cargo.toml b/Cargo.toml index f67b0fcb..50bbee92 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,8 @@ members = [ "node-hub/dora-kit-car", "node-hub/dora-object-to-pose", "node-hub/dora-mistral-rs", + "node-hub/dora-av1-encoder", + "node-hub/dora-dav1d", "libraries/extensions/ros2-bridge", "libraries/extensions/ros2-bridge/msg-gen", "libraries/extensions/ros2-bridge/python", diff --git a/node-hub/dora-av1-encoder/Cargo.toml b/node-hub/dora-av1-encoder/Cargo.toml new file mode 100644 index 00000000..b0ee3d8a --- /dev/null +++ b/node-hub/dora-av1-encoder/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "dora-av1-encoder" +edition = "2021" +version.workspace = true +description.workspace = true +documentation.workspace = true +license.workspace = true +repository.workspace = true + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +rav1e = { version = "0.7.1", features = ["serialize"] } +dora-node-api = { workspace = true, features = ["tracing"] } +eyre = "0.6.8" diff --git a/node-hub/dora-av1-encoder/example.py b/node-hub/dora-av1-encoder/example.py new file mode 100644 index 00000000..80d54b89 --- /dev/null +++ b/node-hub/dora-av1-encoder/example.py @@ -0,0 +1,22 @@ +import av +import cv2 +import numpy as np + +# Open the AV1 video file +container = av.open("video.av1") + +# Get the video stream +stream = next(s for s in container.streams if s.type == "video") + +# Iterate over the video frames +for frame in container.decode(stream): + # Convert the frame to a NumPy array + img = frame.to_ndarray(format="bgr24") + + # Display the frame + cv2.imshow("Frame", img) + if cv2.waitKey(1) & 0xFF == ord("q"): + break + +# Release the resources +cv2.destroyAllWindows() diff --git a/node-hub/dora-av1-encoder/src/main.rs b/node-hub/dora-av1-encoder/src/main.rs new file mode 100644 index 00000000..dbad450e --- /dev/null +++ b/node-hub/dora-av1-encoder/src/main.rs @@ -0,0 +1,217 @@ +// 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::time::Duration; + +use dora_node_api::arrow::array::UInt8Array; +use dora_node_api::dora_core::config::DataId; +use dora_node_api::{DoraNode, Event, IntoArrow, MetadataParameters, Parameter}; +use eyre::{Context as EyreContext, Result}; +// Encode the same tiny blank frame 30 times +use rav1e::config::RateControlConfig; +use rav1e::config::SpeedSettings; + +use rav1e::*; + +fn bgr_to_yuv(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: Vec, width: usize, height: usize) -> (Vec, Vec, Vec) { + // 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 enc = EncoderConfig { + width, + height, + speed_settings: SpeedSettings::from_preset(8), + low_latency: true, + ..Default::default() + }; + + let cfg = Config::new() + // .with_rate_control(RateControlConfig::new().with_emit_data(true)) + .with_encoder_config(enc.clone()) + .with_threads(16); + + let mut ctx: Context = cfg.new_context().unwrap(); + + let (mut node, mut events) = + DoraNode::init_from_env().context("Could not initialize dora node")?; + + let mut time = std::time::Instant::now(); + loop { + let buffer = match events.recv() { + Some(Event::Input { id, data, 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" + }; + + if encoding == "bgr8" { + let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); + let buffer: Vec = buffer.values().to_vec(); + buffer + // Transpose values from BGR to RGB + // let buffer: Vec = buffer.chunks(3).flat_map(|x| [x[2], x[1], x[0]]).collect(); + + //un + } else if encoding == "yuv420" { + let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap(); + let buffer: Vec = buffer.values().to_vec(); + buffer + } 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 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."); + continue; + } + } + Some(Event::Error(e)) => { + continue; + } + _ => break, + }; + //let (y, u, v) = bgr_to_yuv(buffer, 640 as usize, 480 as usize); + + let (y, u, v) = get_yuv_planes(buffer, width, height); + let mut f = ctx.new_frame(); + + let xdec = f.planes[0].cfg.xdec; + let stride = (enc.width + xdec) >> xdec; + f.planes[0].copy_from_raw_u8(&y, stride, 1); + let xdec = f.planes[1].cfg.xdec; + let stride = (enc.width + xdec) >> xdec; + f.planes[1].copy_from_raw_u8(&u, stride, 1); + let xdec = f.planes[2].cfg.xdec; + let stride = (enc.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 => { + println!("Unable to append frame to the internal queue"); + panic!("Unable to send frame "); + } + _ => { + panic!("Unable to send frame "); + } + }, + } + println!("Frame sent to encoder"); + for _ in 0..1 { + match ctx.receive_packet() { + Ok(pkt) => { + println!("Time to encode: {:?}", time.elapsed()); + time = std::time::Instant::now(); + let data = pkt.data; + println!("frame compression: {:#?}", width * height * 3 / data.len()); + println!("frame size: {:#?}", data.len()); + let arrow = data.into_arrow(); + node.send_output( + DataId::from("frame".to_owned()), + MetadataParameters::default(), + arrow, + ) + .context("could not send output") + .unwrap(); + + break; + } + Err(e) => match e { + EncoderStatus::LimitReached => { + break; + } + EncoderStatus::Encoded => { + break; + } + EncoderStatus::NeedMoreData => { + break; + } + _ => { + panic!("Unable to receive packet",); + } + }, + } + } + } + + Ok(()) +} diff --git a/node-hub/dora-dav1d/Cargo.toml b/node-hub/dora-dav1d/Cargo.toml new file mode 100644 index 00000000..a0940226 --- /dev/null +++ b/node-hub/dora-dav1d/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "dora-dav1d" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +dav1d = "0.10" +bitstream-io = "2.0" +log = "0.4" +structopt = "0.3" +dora-node-api = { workspace = true, features = ["tracing"] } +eyre = "0.6.8" diff --git a/node-hub/dora-dav1d/dataflow.yml b/node-hub/dora-dav1d/dataflow.yml new file mode 100644 index 00000000..b4853804 --- /dev/null +++ b/node-hub/dora-dav1d/dataflow.yml @@ -0,0 +1,53 @@ +nodes: + - id: camera + build: pip install ../../node-hub/opencv-video-capture + path: opencv-video-capture + inputs: + tick: dora/timer/millis/50 + outputs: + - image + env: + CAPTURE_PATH: 0 + IMAGE_WIDTH: 640 + IMAGE_HEIGHT: 480 + ENCODING: yuv420 + + - id: av + path: /home/peter/Documents/work/dora/target/release/dora-av1-encoder + build: cargo build -p dora-av1-encoder --release + inputs: + image: camera/image + outputs: + - frame + env: + IMAGE_WIDTH: 640 + IMAGE_HEIGHT: 480 + + - id: dav1d + path: /home/peter/Documents/work/dora/target/release/dora-dav1d + build: cargo build -p dora-dav1d --release + inputs: + image: av/frame + outputs: + - frame + env: + IMAGE_WIDTH: 640 + IMAGE_HEIGHT: 480 + + - id: plot + build: pip install -e ../../node-hub/opencv-plot + path: opencv-plot + inputs: + image: dav1d/frame + env: + IMAGE_WIDTH: 640 + IMAGE_HEIGHT: 480 + + - id: plot2 + build: pip install -e ../../node-hub/opencv-plot + path: opencv-plot + inputs: + image: camera/image + env: + IMAGE_WIDTH: 640 + IMAGE_HEIGHT: 480 diff --git a/node-hub/dora-dav1d/src/main.rs b/node-hub/dora-dav1d/src/main.rs new file mode 100644 index 00000000..0dd758b5 --- /dev/null +++ b/node-hub/dora-dav1d/src/main.rs @@ -0,0 +1,189 @@ +use std::time::Duration; + +use dav1d::Settings; +use dora_node_api::{ + arrow::array::UInt8Array, dora_core::config::DataId, DoraNode, Event, IntoArrow, + MetadataParameters, +}; +use eyre::{Context, Result}; + +fn yuv420_to_bgr( + y_plane: &[u8], + u_plane: &[u8], + v_plane: &[u8], + width: usize, + height: usize, +) -> Vec { + 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); + + let height: usize = std::env::var("IMAGE_HEIGHT") + .unwrap_or_else(|_| "480".to_string()) + .parse() + .unwrap(); + let width: usize = std::env::var("IMAGE_WIDTH") + .unwrap_or_else(|_| "640".to_string()) + .parse() + .unwrap(); + + let mut dec = + dav1d::Decoder::with_settings(&settings).expect("failed to create decoder instance"); + + let (tx, rx) = std::sync::mpsc::sync_channel(1); + let (mut node, mut events) = + DoraNode::init_from_env().context("Could not initialize dora node")?; + + let _ = std::thread::spawn(move || { + let mut now = std::time::Instant::now(); + loop { + while let Ok(data) = rx.recv_timeout(Duration::from_millis(100)) { + match dec.send_data(data, None, None, None) { + Err(e) if e.is_again() => { + panic!("Error sending data to the decoder: {}", e); + if let Ok(p) = dec.get_picture() { + let mut y = p.plane(dav1d::PlanarImageComponent::Y).to_vec(); + let mut u = p.plane(dav1d::PlanarImageComponent::U).to_vec(); + let mut v = p.plane(dav1d::PlanarImageComponent::V).to_vec(); + y.append(&mut u); + y.append(&mut v); + // let rgb = yuv420_to_rgb(&y, &u, &v, 100, 100); + let arrow = y.into_arrow(); + let mut metadata = MetadataParameters::default(); + metadata.insert( + "width".to_string(), + dora_node_api::Parameter::Integer(640), + ); + metadata.insert( + "height".to_string(), + dora_node_api::Parameter::Integer(480), + ); + metadata.insert( + "encoding".to_string(), + dora_node_api::Parameter::String("yuv420".to_string()), + ); + node.send_output(DataId::from("frame".to_string()), metadata, arrow) + .unwrap(); + println!("Time to decode: {:?}", now.elapsed()); + now = std::time::Instant::now(); + } + // If the decoder did not consume all data, output all + // pending pictures and send pending data to the decoder + // until it is all used up. + //loop { + // handle_pending_pictures(&mut dec, false, &mut node); + + // match dec.send_pending_data() { + // Err(e) if e.is_again() => continue, + // Err(e) => { + // panic!("Error sending pending data to the decoder: {}", e); + // } + // _ => break, + // } + //} + } + Err(e) => { + panic!("Error sending data to the decoder: {}", e); + } + Ok(()) => { + if let Ok(p) = dec.get_picture() { + let mut y = p.plane(dav1d::PlanarImageComponent::Y).to_vec(); + let mut u = p.plane(dav1d::PlanarImageComponent::U).to_vec(); + let mut v = p.plane(dav1d::PlanarImageComponent::V).to_vec(); + // u.iter_mut().for_each(|e| { + // if *e < 128 { + // *e = *e + 128 + // } + // }); + // v.iter_mut().for_each(|e: &mut u8| { + // if *e < 128 { + // *e = *e + 128 + // } + // }); + + // y.append(&mut u); + // y.append(&mut v); + let y = yuv420_to_bgr(&y, &u, &v, width, height); + + let arrow = y.into_arrow(); + let mut metadata = MetadataParameters::default(); + metadata.insert( + "width".to_string(), + dora_node_api::Parameter::Integer( + width.try_into().unwrap_or_default(), + ), + ); + metadata.insert( + "height".to_string(), + dora_node_api::Parameter::Integer( + height.try_into().unwrap_or_default(), + ), + ); + metadata.insert( + "encoding".to_string(), + dora_node_api::Parameter::String("bgr8".to_string()), + ); + node.send_output(DataId::from("frame".to_string()), metadata, arrow) + .unwrap(); + println!("Time to decode: {:?}", now.elapsed()); + println!("Delay: {:#?}", dec.get_frame_delay()); + + now = std::time::Instant::now(); + } + } + } + } + } + }); + loop { + match events.recv() { + Some(Event::Input { + id: _, + data, + metadata: _, + }) => { + let data = data.as_any().downcast_ref::().unwrap(); + let data = data.values().clone(); + + // Send packet to the decoder + + tx.try_send(data).unwrap_or_default(); + // Handle all pending pictures before sending the next data. + // handle_pending_pictures(&mut dec, false, &mut node); + } + None => break, + Some(_) => break, + } + } + // Handle all pending pictures that were not output yet. + // handle_pending_pictures(&mut dec, true, &mut node); + + Ok(()) +}