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

Push dav1d and rav1e on pypi

tags/v0.3.11^2
haixuantao 9 months ago
parent
56380d181e
commit
53d6b61839
9 changed files with 599 additions and 456 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +2
    -0
      Cargo.lock
  2. +16
    -0
      node-hub/dora-dav1d/Cargo.toml
  3. +30
    -0
      node-hub/dora-dav1d/pyproject.toml
  4. +129
    -0
      node-hub/dora-dav1d/src/lib.rs
  5. +2
    -107
      node-hub/dora-dav1d/src/main.rs
  6. +17
    -0
      node-hub/dora-rav1e/Cargo.toml
  7. +30
    -0
      node-hub/dora-rav1e/pyproject.toml
  8. +371
    -0
      node-hub/dora-rav1e/src/lib.rs
  9. +2
    -349
      node-hub/dora-rav1e/src/main.rs

+ 2
- 0
Cargo.lock View File

@@ -3183,6 +3183,7 @@ dependencies = [
"dora-node-api",
"eyre",
"log",
"pyo3",
"structopt",
]

@@ -3441,6 +3442,7 @@ dependencies = [
"dora-node-api",
"eyre",
"log",
"pyo3",
"rav1e",
]



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

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

+ 30
- 0
node-hub/dora-dav1d/pyproject.toml View File

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

+ 129
- 0
node-hub/dora-dav1d/src/lib.rs View File

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

#[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(())
}

+ 2
- 107
node-hub/dora-dav1d/src/main.rs View File

@@ -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<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 {
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::<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(())
fn main() -> eyre::Result<()> {
dora_dav1d::lib_main()
}

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

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


+ 30
- 0
node-hub/dora-rav1e/pyproject.toml View File

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

+ 371
- 0
node-hub/dora-rav1e/src/lib.rs View File

@@ -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<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: &[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<u8> = buffer.values().to_vec();
let (y, u, v) = bgr8_to_yuv420(buffer, width, height);

// Transpose values from BGR to RGB
// let buffer: Vec<u8> = buffer.chunks(3).flat_map(|x| [x[2], x[1], x[0]]).collect();

let mut ctx: Context<u8> = 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<u8> = 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<u16> = 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<u8> = buffer.values().to_vec();
let buffer: Vec<u8> =
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<u8> = 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(())
}

+ 2
- 349
node-hub/dora-rav1e/src/main.rs View File

@@ -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<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: &[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<u8> = buffer.values().to_vec();
let (y, u, v) = bgr8_to_yuv420(buffer, width, height);

// Transpose values from BGR to RGB
// let buffer: Vec<u8> = buffer.chunks(3).flat_map(|x| [x[2], x[1], x[0]]).collect();

let mut ctx: Context<u8> = 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<u8> = 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<u16> = 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<u8> = buffer.values().to_vec();
let buffer: Vec<u8> =
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<u8> = 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()
}

Write Preview
Loading…
Cancel
Save