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

adding avif support within rav1e

tags/v0.3.12-rc0
haixuantao haixuanTao 9 months ago
parent
489d777203
commit
70074749ba
4 changed files with 233 additions and 158 deletions
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. +13
    -0
      node-hub/dora-dav1d/src/lib.rs
  2. +1
    -0
      node-hub/dora-rav1e/Cargo.toml
  3. +218
    -157
      node-hub/dora-rav1e/src/lib.rs
  4. +1
    -1
      node-hub/dora-rerun/src/lib.rs

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

@@ -58,6 +58,19 @@ pub fn lib_main() -> Result<()> {
}) => {
if let Some(data) = data.as_any().downcast_ref::<UInt8Array>() {
let data = data.values().clone();
let encoding = metadata
.parameters
.get("encoding")
.and_then(|p| match p {
dora_node_api::Parameter::String(s) => Some(s),
_ => None,
})
.map(|s| s.as_str())
.unwrap_or("av1");
if encoding != "av1" {
warn!("Unsupported encoding {}", encoding);
continue;
}
match dec.send_data(data, None, None, None) {
Err(e) => {
warn!("Error sending data to the decoder: {}", e);


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

@@ -25,6 +25,7 @@ pyo3 = { workspace = true, features = [
"eyre",
"generate-import-lib",
], optional = true }
avif-serialize = "0.8.3"


[lib]


+ 218
- 157
node-hub/dora-rav1e/src/lib.rs View File

@@ -10,14 +10,47 @@
use std::env::var;

use dora_node_api::arrow::array::{UInt16Array, UInt8Array};
use dora_node_api::{DoraNode, Event, IntoArrow, Parameter};
use dora_node_api::dora_core::config::DataId;
use dora_node_api::{DoraNode, Event, IntoArrow, Metadata, Parameter};
use eyre::{Context as EyreContext, Result};
use log::warn;
use rav1e::color::{ColorDescription, MatrixCoefficients};
// Encode the same tiny blank frame 30 times
use rav1e::config::SpeedSettings;

use rav1e::*;

pub fn fill_zeros_toward_center_y_plane_in_place(y: &mut [u16], width: usize, height: usize) {
assert_eq!(y.len(), width * height);

for row in 0..height {
let row_start = row * width;
let center = width / 2;

// --- Fill left half (left to center)
let mut last = 0u16;
for col in 0..center {
let idx = row_start + col;
if y[idx] == 0 {
y[idx] = last;
} else {
last = y[idx];
}
}

// --- Fill right half (right to center)
last = 0u16;
for col in (center..width).rev() {
let idx = row_start + col;
if y[idx] == 0 {
y[idx] = last;
} else {
last = y[idx];
}
}
}
}

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)];
@@ -69,6 +102,123 @@ fn get_yuv_planes(buffer: &[u8], width: usize, height: usize) -> (&[u8], &[u8],
(y_plane, u_plane, v_plane)
}

fn send_yuv(
y: &[u8],
u: &[u8],
v: &[u8],
enc: EncoderConfig,
width: usize,
height: usize,
node: &mut DoraNode,
id: DataId,
metadata: &mut Metadata,
output_enconding: &str,
) -> () {
// Create a new Arrow array for the YUV420 data
let cfg = Config::new().with_encoder_config(enc.clone());
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 ");
}
},
}
ctx.flush();
match ctx.receive_packet() {
Ok(pkt) => match output_enconding {
"avif" => {
let data = pkt.data.clone();
metadata.parameters.insert(
"encoding".to_string(),
Parameter::String("avif".to_string()),
);
let matrix_coefficients = if let Some(desc) = enc.color_description {
desc.matrix_coefficients
} else {
MatrixCoefficients::BT709
};
let data = avif_serialize::Aviffy::new()
.set_chroma_subsampling((true, true))
.set_seq_profile(0)
.matrix_coefficients(match matrix_coefficients {
MatrixCoefficients::Identity => {
avif_serialize::constants::MatrixCoefficients::Rgb
}
MatrixCoefficients::BT709 => {
avif_serialize::constants::MatrixCoefficients::Bt709
}
MatrixCoefficients::Unspecified => {
avif_serialize::constants::MatrixCoefficients::Unspecified
}
MatrixCoefficients::BT601 => {
avif_serialize::constants::MatrixCoefficients::Bt601
}
MatrixCoefficients::YCgCo => {
avif_serialize::constants::MatrixCoefficients::Ycgco
}
MatrixCoefficients::BT2020NCL => {
avif_serialize::constants::MatrixCoefficients::Bt2020Ncl
}
MatrixCoefficients::BT2020CL => {
avif_serialize::constants::MatrixCoefficients::Bt2020Cl
}
_ => {
warn!("color matrix coefficients");
avif_serialize::constants::MatrixCoefficients::Rgb
}
})
.to_vec(
&data,
None,
width as u32,
height as u32,
enc.bit_depth as u8,
);

let arrow = data.into_arrow();
node.send_output(id, metadata.parameters.clone(), arrow)
.context("could not send output")
.unwrap();
}
_ => {
metadata
.parameters
.insert("encoding".to_string(), Parameter::String("av1".to_string()));
let data = pkt.data;
let arrow = data.into_arrow();
node.send_output(id, metadata.parameters.clone(), arrow)
.context("could not send output")
.unwrap();
}
},
Err(e) => match e {
EncoderStatus::LimitReached => {}
EncoderStatus::Encoded => {}
EncoderStatus::NeedMoreData => {}
_ => {
panic!("Unable to receive packet",);
}
},
}
}

pub fn lib_main() -> Result<()> {
let mut height = std::env::var("IMAGE_HEIGHT")
.unwrap_or_else(|_| "480".to_string())
@@ -79,10 +229,17 @@ pub fn lib_main() -> Result<()> {
.parse()
.unwrap();
let speed = var("RAV1E_SPEED").map(|s| s.parse().unwrap()).unwrap_or(10);
let output_encoding = var("ENCODING").unwrap_or("av1".to_string());
let mut enc = EncoderConfig {
width,
height,
speed_settings: SpeedSettings::from_preset(speed),
chroma_sampling: color::ChromaSampling::Cs420,
color_description: Some(ColorDescription {
matrix_coefficients: MatrixCoefficients::BT709,
transfer_characteristics: color::TransferCharacteristics::BT709,
color_primaries: color::ColorPrimaries::BT709,
}),
low_latency: true,
..Default::default()
};
@@ -127,115 +284,49 @@ pub fn lib_main() -> Result<()> {
_ => {}
}

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",);
}
},
}
send_yuv(
&y,
&u,
&v,
enc,
width,
height,
&mut node,
id,
&mut metadata,
&output_encoding,
);
} 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" {
send_yuv(
&y,
&u,
&v,
enc,
width,
height,
&mut node,
id,
&mut metadata,
&output_encoding,
);
} else if encoding == "mono16" || encoding == "z16" {
let buffer: &UInt16Array = data.as_any().downcast_ref().unwrap();
let buffer: &[u16] = buffer.values();
let mut buffer = buffer.values().to_vec();
if std::env::var("FILL_ZEROS")
.map(|s| s != "false")
.unwrap_or(true)
{
fill_zeros_toward_center_y_plane_in_place(&mut buffer, width, height);
}

// let buffer = shift_u16_slice_to_upper_12_bits(buffer);
let bytes: &[u8] = &bytemuck::cast_slice(&buffer);

@@ -258,17 +349,25 @@ pub fn lib_main() -> Result<()> {
}
},
}
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();
match output_encoding.as_str() {
"avif" => {
warn!("avif encoding not supported for mono16");
}
_ => {
metadata.parameters.insert(
"encoding".to_string(),
Parameter::String("av1".to_string()),
);
let arrow = data.into_arrow();
node.send_output(id, metadata.parameters, arrow)
.context("could not send output")
.unwrap();
}
}
}
Err(e) => match e {
EncoderStatus::LimitReached => {}
@@ -285,56 +384,18 @@ pub fn lib_main() -> Result<()> {
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",);
}
},
}
send_yuv(
&y,
&u,
&v,
enc,
width,
height,
&mut node,
id,
&mut metadata,
&output_encoding,
);
} else {
unimplemented!("We haven't worked on additional encodings.");
}


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

@@ -155,7 +155,7 @@ pub fn lib_main() -> Result<()> {
);
rec.log(id.as_str(), &image)
.context("could not log image")?;
} else if ["jpeg", "png"].contains(&encoding) {
} else if ["jpeg", "png", "avif"].contains(&encoding) {
let buffer: &UInt8Array = data.as_any().downcast_ref().unwrap();
let buffer: &[u8] = buffer.values();



Write Preview
Loading…
Cancel
Save