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

Initial support for dav1d and rav1e

tags/v0.3.11-rc1
haixuanTao haixuantao 1 year ago
parent
ced06a9b54
commit
a82c2a451b
8 changed files with 784 additions and 2 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +273
    -2
      Cargo.lock
  2. +2
    -0
      Cargo.toml
  3. +15
    -0
      node-hub/dora-av1-encoder/Cargo.toml
  4. +22
    -0
      node-hub/dora-av1-encoder/example.py
  5. +217
    -0
      node-hub/dora-av1-encoder/src/main.rs
  6. +13
    -0
      node-hub/dora-dav1d/Cargo.toml
  7. +53
    -0
      node-hub/dora-dav1d/dataflow.yml
  8. +189
    -0
      node-hub/dora-dav1d/src/main.rs

+ 273
- 2
Cargo.lock View File

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


+ 2
- 0
Cargo.toml View File

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


+ 15
- 0
node-hub/dora-av1-encoder/Cargo.toml View File

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

+ 22
- 0
node-hub/dora-av1-encoder/example.py View File

@@ -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()

+ 217
- 0
node-hub/dora-av1-encoder/src/main.rs View File

@@ -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<u8>, width: usize, height: usize) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
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<u8>, width: usize, height: usize) -> (Vec<u8>, Vec<u8>, Vec<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 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<u16> = 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<u8> = buffer.values().to_vec();
buffer
// Transpose values from BGR to RGB
// let buffer: Vec<u8> = 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<u8> = 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(())
}

+ 13
- 0
node-hub/dora-dav1d/Cargo.toml View File

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

+ 53
- 0
node-hub/dora-dav1d/dataflow.yml View File

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

+ 189
- 0
node-hub/dora-dav1d/src/main.rs View File

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

Write Preview
Loading…
Cancel
Save