diff --git a/Cargo.lock b/Cargo.lock index 1d5a663..22b0c82 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -33,6 +33,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + [[package]] name = "anstream" version = "0.6.21" @@ -125,6 +140,13 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "av1_parser" +version = "0.1.0" +dependencies = [ + "anyhow", +] + [[package]] name = "bitflags" version = "1.3.2" @@ -196,6 +218,12 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.55" @@ -212,6 +240,33 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "clap" version = "4.5.57" @@ -363,6 +418,41 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools 0.13.0", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools 0.13.0", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -388,6 +478,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "darling" version = "0.20.11" @@ -513,6 +609,8 @@ dependencies = [ "bitvec", "bitvec_helpers", "crc", + "criterion", + "libc", "roxmltree", "serde", "serde_json", @@ -526,6 +624,7 @@ dependencies = [ "anyhow", "assert_cmd", "assert_fs", + "av1_parser", "bitvec_helpers", "clap", "clap_lex", @@ -533,7 +632,7 @@ dependencies = [ "hdr10plus", "hevc_parser", "indicatif", - "itertools", + "itertools 0.14.0", "madvr_parse", "plotters", "predicates", @@ -750,6 +849,17 @@ dependencies = [ "walkdir", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.16.1" @@ -852,6 +962,15 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -1010,6 +1129,12 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "option-ext" version = "0.2.0" @@ -1025,6 +1150,16 @@ dependencies = [ "ttf-parser 0.25.1", ] +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "pathfinder_geometry" version = "0.5.1" @@ -1064,6 +1199,7 @@ dependencies = [ "pathfinder_geometry", "plotters-backend", "plotters-bitmap", + "plotters-svg", "ttf-parser 0.20.0", "wasm-bindgen", "web-sys", @@ -1085,6 +1221,15 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "png" version = "0.17.16" @@ -1170,6 +1315,26 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_users" version = "0.5.2" @@ -1429,6 +1594,16 @@ dependencies = [ "time-core", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.10.0" @@ -1766,6 +1941,26 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "zerocopy" +version = "0.8.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2578b716f8a7a858b7f02d5bd870c14bf4ddbbcf3a4c05414ba6503640505e3" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e6cc098ea4d3bd6246687de65af3f920c430e236bee1e3bf2e441463f08a02f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zmij" version = "1.0.19" diff --git a/Cargo.toml b/Cargo.toml index cdf4eb4..4d9f213 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,7 @@ +[workspace] +members = [".", "dolby_vision", "av1_parser"] +resolver = "2" + [package] name = "dovi_tool" version = "2.3.1" @@ -18,6 +22,7 @@ bitvec_helpers = { version = "4.0.1", default-features = false, features = ["bit hevc_parser = { version = "0.6.10", features = ["hevc_io"] } madvr_parse = "1.0.3" hdr10plus = { version = "2.1.5", features = ["json"] } +av1_parser = { path = "av1_parser" } anyhow = "1.0.101" clap = { version = "4.5.57", features = ["derive", "wrap_help", "deprecated"] } diff --git a/av1_parser/Cargo.toml b/av1_parser/Cargo.toml new file mode 100644 index 0000000..fc69e6c --- /dev/null +++ b/av1_parser/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "av1_parser" +version = "0.1.0" +edition = "2024" + +[dependencies] +anyhow = "1.0" diff --git a/av1_parser/src/lib.rs b/av1_parser/src/lib.rs new file mode 100644 index 0000000..a9d3da4 --- /dev/null +++ b/av1_parser/src/lib.rs @@ -0,0 +1,299 @@ +#![allow(dead_code)] + +use std::io::{BufRead, ErrorKind, Read, Write}; + +use anyhow::{Result, bail}; + +// --------------------------------------------------------------------------- +// OBU type constants (AV1 spec Table 5) +// --------------------------------------------------------------------------- +pub const OBU_SEQUENCE_HEADER: u8 = 1; +pub const OBU_TEMPORAL_DELIMITER: u8 = 2; +pub const OBU_FRAME_HEADER: u8 = 3; +pub const OBU_METADATA: u8 = 5; +pub const OBU_FRAME: u8 = 6; +pub const OBU_REDUNDANT_FRAME_HEADER: u8 = 7; + +// --------------------------------------------------------------------------- +// Obu — a single parsed OBU with its complete raw bytes +// --------------------------------------------------------------------------- + +/// A single parsed AV1 Open Bitstream Unit. +pub struct Obu { + pub obu_type: u8, + pub temporal_id: u8, + pub spatial_id: u8, + /// Decoded payload bytes (after header + LEB128 size). + pub payload: Vec, + /// Complete raw bytes of this OBU as it appeared on disk. + /// Used for pass-through writing. + pub raw_bytes: Vec, +} + +impl Obu { + /// Read one OBU from `reader`. Returns `None` on clean EOF. + /// + /// Only supports the *Low Overhead Bitstream Format* where every OBU + /// carries a size field (`obu_has_size_field == 1`). + pub fn read_from(reader: &mut R) -> Result> { + // ---- header byte ---- + let mut header_byte = [0u8; 1]; + match reader.read_exact(&mut header_byte) { + Ok(()) => {} + Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(None), + Err(e) => return Err(e.into()), + } + + let byte = header_byte[0]; + if byte >> 7 != 0 { + bail!("AV1 OBU forbidden bit is set (byte = 0x{byte:02X})"); + } + + let obu_type = (byte >> 3) & 0x0F; + let has_extension = (byte >> 2) & 1 != 0; + let has_size_field = (byte >> 1) & 1 != 0; + + let mut raw = vec![byte]; + let mut temporal_id = 0u8; + let mut spatial_id = 0u8; + + // ---- optional extension header ---- + if has_extension { + let mut ext = [0u8; 1]; + reader.read_exact(&mut ext)?; + temporal_id = (ext[0] >> 5) & 0x07; + spatial_id = (ext[0] >> 3) & 0x03; + raw.push(ext[0]); + } + + if !has_size_field { + bail!( + "OBU (type {obu_type}) has no size field; \ + only Low Overhead Bitstream Format is supported" + ); + } + + // ---- LEB128 payload size ---- + let payload_size = { + let mut size: u64 = 0; + let mut shift = 0u32; + loop { + let mut b = [0u8; 1]; + reader.read_exact(&mut b)?; + raw.push(b[0]); + size |= ((b[0] & 0x7F) as u64) << shift; + shift += 7; + if b[0] & 0x80 == 0 { + break; + } + if shift >= 56 { + bail!("LEB128 overflow while reading OBU size"); + } + } + size as usize + }; + + // ---- payload ---- + let payload_start = raw.len(); + raw.resize(payload_start + payload_size, 0); + reader.read_exact(&mut raw[payload_start..])?; + let payload = raw[payload_start..].to_vec(); + + Ok(Some(Obu { + obu_type, + temporal_id, + spatial_id, + payload, + raw_bytes: raw, + })) + } +} + +// --------------------------------------------------------------------------- +// LEB128 encoding / decoding +// --------------------------------------------------------------------------- + +/// Encode a `u64` value as LEB128 (unsigned). +pub fn encode_leb128(mut value: u64) -> Vec { + let mut result = Vec::new(); + loop { + let mut byte = (value & 0x7F) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + result.push(byte); + if value == 0 { + break; + } + } + result +} + +/// Decode a LEB128-encoded value from `data`. +/// Returns `(value, bytes_consumed)`. +pub fn decode_leb128(data: &[u8]) -> (u64, usize) { + let mut value = 0u64; + let mut bytes_read = 0usize; + for (i, &byte) in data.iter().enumerate() { + if i >= 8 { + break; + } + value |= ((byte & 0x7F) as u64) << (7 * i); + bytes_read += 1; + if byte & 0x80 == 0 { + break; + } + } + (value, bytes_read) +} + +// --------------------------------------------------------------------------- +// IVF container support +// --------------------------------------------------------------------------- + +/// IVF file signature ("DKIF"). +pub const IVF_SIGNATURE: [u8; 4] = *b"DKIF"; + +/// Size of the IVF file header in bytes. +pub const IVF_FILE_HEADER_LEN: usize = 32; + +/// Size of an IVF frame header in bytes. +pub const IVF_FRAME_HEADER_LEN: usize = 12; + +/// Header of a single IVF frame. +pub struct IvfFrameHeader { + /// Number of bytes in the frame data that follows. + pub frame_size: u32, + /// Presentation timestamp (in stream timebase). + pub timestamp: u64, +} + +/// Probe the first bytes of `reader` to decide whether the stream is an IVF +/// container. If the IVF signature is detected the 32-byte file header is +/// consumed from `reader` and returned; otherwise `None` is returned and +/// **no bytes are consumed**. +pub fn try_read_ivf_file_header( + reader: &mut R, +) -> Result> { + { + let buf = reader.fill_buf()?; + if buf.len() < 4 || buf[..4] != IVF_SIGNATURE { + return Ok(None); + } + } + let mut header = [0u8; IVF_FILE_HEADER_LEN]; + reader.read_exact(&mut header)?; + Ok(Some(header)) +} + +/// Read one IVF frame header from `reader`. Returns `None` on clean EOF. +pub fn read_ivf_frame_header(reader: &mut R) -> Result> { + let mut buf = [0u8; IVF_FRAME_HEADER_LEN]; + match reader.read_exact(&mut buf) { + Ok(()) => {} + Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(None), + Err(e) => return Err(e.into()), + } + Ok(Some(IvfFrameHeader { + frame_size: u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]), + timestamp: u64::from_le_bytes(buf[4..12].try_into().unwrap()), + })) +} + +/// Write an IVF frame header (frame_size + timestamp) to `writer`. +pub fn write_ivf_frame_header( + writer: &mut W, + frame_size: u32, + timestamp: u64, +) -> Result<()> { + writer.write_all(&frame_size.to_le_bytes())?; + writer.write_all(×tamp.to_le_bytes())?; + Ok(()) +} + +/// Read all OBUs from a single IVF frame's data bytes. +pub fn read_obus_from_ivf_frame(frame_data: Vec) -> Result> { + let mut cursor = std::io::Cursor::new(frame_data); + let mut obus = Vec::new(); + while let Some(obu) = Obu::read_from(&mut cursor)? { + obus.push(obu); + } + Ok(obus) +} + +// --------------------------------------------------------------------------- +// I/O structs +// --------------------------------------------------------------------------- + +/// Iterates OBUs from a raw AV1 byte stream. +pub struct ObuReader { + reader: R, +} + +impl ObuReader { + pub fn new(reader: R) -> Self { + ObuReader { reader } + } + pub fn next_obu(&mut self) -> Result> { + Obu::read_from(&mut self.reader) + } + pub fn into_inner(self) -> R { + self.reader + } +} + +impl Iterator for ObuReader { + type Item = Result; + fn next(&mut self) -> Option { + self.next_obu().transpose() + } +} + +/// Writes IVF frames. Writes the file header in `new()`. +pub struct IvfWriter { + writer: W, +} + +impl IvfWriter { + /// Writes the 32-byte IVF file header immediately. + pub fn new(mut writer: W, file_header: &[u8; 32]) -> Result { + writer.write_all(file_header)?; + Ok(IvfWriter { writer }) + } + + /// Writes one IVF frame (12-byte frame header + frame data). + pub fn write_frame(&mut self, timestamp: u64, frame_data: &[u8]) -> Result<()> { + write_ivf_frame_header(&mut self.writer, frame_data.len() as u32, timestamp)?; + self.writer.write_all(frame_data)?; + Ok(()) + } + + pub fn flush(&mut self) -> Result<()> { + self.writer.flush().map_err(Into::into) + } + + pub fn into_inner(self) -> W { + self.writer + } +} + +/// Writes raw AV1 OBUs directly. +pub struct ObuWriter { + writer: W, +} + +impl ObuWriter { + pub fn new(writer: W) -> Self { + ObuWriter { writer } + } + pub fn write_raw(&mut self, bytes: &[u8]) -> Result<()> { + self.writer.write_all(bytes).map_err(Into::into) + } + pub fn flush(&mut self) -> Result<()> { + self.writer.flush().map_err(Into::into) + } + pub fn into_inner(self) -> W { + self.writer + } +} diff --git a/src/commands/convert.rs b/src/commands/convert.rs index ea18de0..8de76f0 100644 --- a/src/commands/convert.rs +++ b/src/commands/convert.rs @@ -5,7 +5,7 @@ use std::path::PathBuf; pub struct ConvertArgs { #[arg( id = "input", - help = "Sets the input HEVC file to use, or piped with -", + help = "Sets the input file to use (.hevc, .av1, .ivf), or piped with -", long, short = 'i', conflicts_with = "input_pos", @@ -16,7 +16,7 @@ pub struct ConvertArgs { #[arg( id = "input_pos", - help = "Sets the input HEVC file to use, or piped with - (positional)", + help = "Sets the input file to use (.hevc, .av1, .ivf), or piped with - (positional)", conflicts_with = "input", required_unless_present = "input", value_hint = ValueHint::FilePath diff --git a/src/commands/extract_rpu.rs b/src/commands/extract_rpu.rs index 793a0f3..97fd519 100644 --- a/src/commands/extract_rpu.rs +++ b/src/commands/extract_rpu.rs @@ -5,7 +5,7 @@ use std::path::PathBuf; pub struct ExtractRpuArgs { #[arg( id = "input", - help = "Sets the input HEVC file to use, or piped with -", + help = "Sets the input file to use (.hevc, .av1, .ivf), or piped with -", long, short = 'i', conflicts_with = "input_pos", @@ -16,7 +16,7 @@ pub struct ExtractRpuArgs { #[arg( id = "input_pos", - help = "Sets the input HEVC file to use, or piped with - (positional)", + help = "Sets the input file to use (.hevc, .av1, .ivf), or piped with - (positional)", conflicts_with = "input", required_unless_present = "input", value_hint = ValueHint::FilePath diff --git a/src/commands/inject_rpu.rs b/src/commands/inject_rpu.rs index ca20521..073dc78 100644 --- a/src/commands/inject_rpu.rs +++ b/src/commands/inject_rpu.rs @@ -5,7 +5,7 @@ use std::path::PathBuf; pub struct InjectRpuArgs { #[arg( id = "input", - help = "Sets the input HEVC file to use", + help = "Sets the input file to use (.hevc, .av1, .ivf)", long, short = 'i', conflicts_with = "input_pos", @@ -16,7 +16,7 @@ pub struct InjectRpuArgs { #[arg( id = "input_pos", - help = "Sets the input HEVC file to use (positional)", + help = "Sets the input file to use (.hevc, .av1, .ivf) (positional)", conflicts_with = "input", required_unless_present = "input", value_hint = ValueHint::FilePath @@ -34,6 +34,6 @@ pub struct InjectRpuArgs { )] pub output: Option, - #[arg(long, num_args = 0, help = "Disable adding AUD NALUs between frames")] + #[arg(long, num_args = 0, help = "Disable adding AUD NALUs between frames (HEVC only)")] pub no_add_aud: bool, } diff --git a/src/commands/remove.rs b/src/commands/remove.rs index a89dadb..826d73e 100644 --- a/src/commands/remove.rs +++ b/src/commands/remove.rs @@ -5,7 +5,7 @@ use std::path::PathBuf; pub struct RemoveArgs { #[arg( id = "input", - help = "Sets the input HEVC file to use, or piped with -", + help = "Sets the input file to use (.hevc, .av1, .ivf), or piped with -", long, short = 'i', conflicts_with = "input_pos", @@ -16,7 +16,7 @@ pub struct RemoveArgs { #[arg( id = "input_pos", - help = "Sets the input HEVC file to use, or piped with - (positional)", + help = "Sets the input file to use (.hevc, .av1, .ivf), or piped with - (positional)", conflicts_with = "input", required_unless_present = "input", value_hint = ValueHint::FilePath diff --git a/src/dovi/av1.rs b/src/dovi/av1.rs new file mode 100644 index 0000000..6f32152 --- /dev/null +++ b/src/dovi/av1.rs @@ -0,0 +1,103 @@ +// Re-export everything the rest of the codebase uses from the av1_parser crate +#[allow(unused_imports)] +pub use av1_parser::{ + IvfFrameHeader, IvfWriter, Obu, ObuReader, ObuWriter, + OBU_TEMPORAL_DELIMITER, OBU_METADATA, + encode_leb128, decode_leb128, + try_read_ivf_file_header, read_ivf_frame_header, write_ivf_frame_header, + read_obus_from_ivf_frame, +}; + +use anyhow::Result; +use dolby_vision::rpu::dovi_rpu::DoviRpu; + +use dolby_vision::av1::ITU_T35_DOVI_RPU_PAYLOAD_HEADER; + +/// Metadata type for ITU-T T.35 +const METADATA_TYPE_ITUT_T35: u64 = 4; + +/// Dolby Vision T.35 country code (United States) +const DOVI_COUNTRY_CODE: u8 = 0xB5; + +/// Returns the T.35 payload bytes (starting at `0xB5` country code) if this +/// `OBU_METADATA` payload contains a Dolby Vision RPU. +/// +/// Layout after `metadata_type = 4` (LEB128): +/// ```text +/// country_code (u8) = 0xB5 +/// terminal_provider_code (u16 BE) = 0x003B +/// terminal_provider_oriented_code (u32 BE) = 0x00000800 +/// +/// ``` +pub fn extract_dovi_t35_payload(obu_payload: &[u8]) -> Option<&[u8]> { + if obu_payload.is_empty() { + return None; + } + + // metadata_type (LEB128) must be 4 + let (mt, mt_len) = decode_leb128(obu_payload); + if mt != METADATA_TYPE_ITUT_T35 { + return None; + } + + let t35 = &obu_payload[mt_len..]; + + // Must start with Dolby Vision country code + if t35.is_empty() || t35[0] != DOVI_COUNTRY_CODE { + return None; + } + + // After country code, the next bytes must match the Dolby Vision header + let after_cc = &t35[1..]; + let hdr_len = ITU_T35_DOVI_RPU_PAYLOAD_HEADER.len(); + if after_cc.len() < hdr_len { + return None; + } + + if &after_cc[..hdr_len] == ITU_T35_DOVI_RPU_PAYLOAD_HEADER { + Some(t35) // return slice starting at 0xB5 + } else { + None + } +} + +/// Returns `true` if this OBU is an `OBU_METADATA` carrying a Dolby Vision RPU. +pub fn is_dovi_rpu_obu(obu: &Obu) -> bool { + obu.obu_type == OBU_METADATA && extract_dovi_t35_payload(&obu.payload).is_some() +} + +/// Build a complete `OBU_METADATA` unit containing the Dolby Vision RPU. +/// +/// Structure: +/// ```text +/// OBU header byte = 0x2A (type=5, has_size_field=1) +/// OBU size (LEB128) +/// metadata_type (LEB128) = 4 +/// 0xB5 country_code +/// +/// ``` +pub fn build_dovi_obu(rpu: &DoviRpu) -> Result> { + // write_av1_rpu_metadata_obu_t35_complete returns: 0xB5 + EMDF payload + let t35_complete = rpu.write_av1_rpu_metadata_obu_t35_complete()?; + + // OBU_METADATA payload: metadata_type(LEB128=4) + T.35 complete payload + let mut obu_payload = encode_leb128(METADATA_TYPE_ITUT_T35); + obu_payload.extend_from_slice(&t35_complete); + + // OBU header byte: + // bit 7: forbidden = 0 + // bits 6-3: obu_type = 5 (OBU_METADATA) + // bit 2: obu_extension_flag = 0 + // bit 1: obu_has_size_field = 1 + // bit 0: reserved = 0 + // => (5 << 3) | 0x02 = 0x2A + let header_byte = (OBU_METADATA << 3) | 0x02u8; + let size_bytes = encode_leb128(obu_payload.len() as u64); + + let mut result = Vec::with_capacity(1 + size_bytes.len() + obu_payload.len()); + result.push(header_byte); + result.extend_from_slice(&size_bytes); + result.extend_from_slice(&obu_payload); + + Ok(result) +} diff --git a/src/dovi/converter.rs b/src/dovi/converter.rs index 4b2cdad..8302fad 100644 --- a/src/dovi/converter.rs +++ b/src/dovi/converter.rs @@ -1,13 +1,28 @@ use anyhow::{Result, bail}; use indicatif::ProgressBar; -use std::path::PathBuf; +use std::fs::File; +use std::io::{BufReader, BufWriter}; +use std::path::{Path, PathBuf}; use crate::commands::ConvertArgs; +use super::av1::{ + IvfWriter, ObuReader, ObuWriter, + build_dovi_obu, is_dovi_rpu_obu, extract_dovi_t35_payload, + try_read_ivf_file_header, read_ivf_frame_header, read_obus_from_ivf_frame, +}; use super::{CliOptions, IoFormat, general_read_write, input_from_either}; +use dolby_vision::rpu::dovi_rpu::DoviRpu; use general_read_write::{DoviProcessor, DoviWriter}; +fn is_av1_input(path: &Path) -> bool { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + pub struct Converter { format: IoFormat, input: PathBuf, @@ -26,16 +41,21 @@ impl Converter { options.discard_el = discard; let input = input_from_either("convert", input, input_pos)?; - let format = hevc_parser::io::format_from_path(&input)?; - let output = match output { - Some(path) => path, - None => match options.discard_el { + let (format, default_output) = if is_av1_input(&input) { + let ext = input.extension().and_then(|e| e.to_str()).unwrap_or("av1"); + (IoFormat::Raw, PathBuf::from(format!("converted.{ext}"))) + } else { + let format = hevc_parser::io::format_from_path(&input)?; + let default = match options.discard_el { true => PathBuf::from("BL_RPU.hevc"), false => PathBuf::from("BL_EL_RPU.hevc"), - }, + }; + (format, default) }; + let output = output.unwrap_or(default_output); + Ok(Self { format, input, @@ -49,6 +69,10 @@ impl Converter { } fn process_input(&self, options: CliOptions) -> Result<()> { + if is_av1_input(&self.input) { + return self.convert_av1(&options); + } + let pb = super::initialize_progress_bar(&self.format, &self.input)?; match self.format { @@ -57,6 +81,70 @@ impl Converter { } } + fn convert_av1(&self, options: &CliOptions) -> Result<()> { + println!("Converting DoVi RPU in AV1 bitstream..."); + + let in_file = File::open(&self.input)?; + let mut reader = BufReader::new(in_file); + + if let Some(ivf_header) = try_read_ivf_file_header(&mut reader)? { + let out_file = BufWriter::new(File::create(&self.output).expect("Can't create file")); + let mut ivf_writer = IvfWriter::new(out_file, &ivf_header)?; + + while let Some(frame_hdr) = read_ivf_frame_header(&mut reader)? { + let mut frame_data = vec![0u8; frame_hdr.frame_size as usize]; + std::io::Read::read_exact(&mut reader, &mut frame_data)?; + + let obus = read_obus_from_ivf_frame(frame_data)?; + let mut new_frame: Vec = Vec::new(); + + for obu in &obus { + if is_dovi_rpu_obu(obu) { + if let Some(t35_payload) = extract_dovi_t35_payload(&obu.payload) { + let mut dovi_rpu = + DoviRpu::parse_itu_t35_dovi_metadata_obu(t35_payload)?; + super::convert_encoded_from_opts_rpu(options, &mut dovi_rpu)?; + let converted_bytes = build_dovi_obu(&dovi_rpu)?; + new_frame.extend_from_slice(&converted_bytes); + } else { + new_frame.extend_from_slice(&obu.raw_bytes); + } + } else { + new_frame.extend_from_slice(&obu.raw_bytes); + } + } + + ivf_writer.write_frame(frame_hdr.timestamp, &new_frame)?; + } + + ivf_writer.flush()?; + } else { + let out_file = BufWriter::new(File::create(&self.output).expect("Can't create file")); + let mut obu_writer = ObuWriter::new(out_file); + let mut obu_reader = ObuReader::new(reader); + + while let Some(obu) = obu_reader.next_obu()? { + if is_dovi_rpu_obu(&obu) { + if let Some(t35_payload) = extract_dovi_t35_payload(&obu.payload) { + let mut dovi_rpu = DoviRpu::parse_itu_t35_dovi_metadata_obu(t35_payload)?; + super::convert_encoded_from_opts_rpu(options, &mut dovi_rpu)?; + let converted_bytes = build_dovi_obu(&dovi_rpu)?; + obu_writer.write_raw(&converted_bytes)?; + } else { + obu_writer.write_raw(&obu.raw_bytes)?; + } + } else { + obu_writer.write_raw(&obu.raw_bytes)?; + } + } + + obu_writer.flush()?; + } + + println!("Done."); + Ok(()) + } + fn convert_raw_hevc(&self, pb: ProgressBar, options: CliOptions) -> Result<()> { let dovi_writer = DoviWriter::new(None, None, None, Some(&self.output)); let mut dovi_processor = DoviProcessor::new( diff --git a/src/dovi/mod.rs b/src/dovi/mod.rs index f198b63..feeb8f2 100644 --- a/src/dovi/mod.rs +++ b/src/dovi/mod.rs @@ -13,6 +13,7 @@ use hevc_parser::io::{IoFormat, StartCodePreset}; use self::editor::EditConfig; use super::commands::ConversionModeCli; +pub mod av1; pub mod converter; pub mod demuxer; pub mod editor; @@ -109,6 +110,24 @@ pub fn convert_encoded_from_opts(opts: &CliOptions, data: &[u8]) -> Result Result<()> { + if let Some(edit_config) = &opts.edit_config { + edit_config.execute_single_rpu(rpu)?; + } else { + if let Some(mode) = opts.mode { + rpu.convert_with_mode(mode)?; + } + + if opts.crop { + rpu.crop()?; + } + } + + Ok(()) +} + pub fn input_from_either(cmd: &str, in1: Option, in2: Option) -> Result { match in1 { Some(in1) => Ok(in1), diff --git a/src/dovi/remover.rs b/src/dovi/remover.rs index 3d4bcb7..320df84 100644 --- a/src/dovi/remover.rs +++ b/src/dovi/remover.rs @@ -1,13 +1,27 @@ use anyhow::{Result, bail}; use indicatif::ProgressBar; -use std::path::PathBuf; +use std::fs::File; +use std::io::{BufReader, BufWriter}; +use std::path::{Path, PathBuf}; use crate::commands::RemoveArgs; +use super::av1::{ + IvfWriter, ObuReader, ObuWriter, + is_dovi_rpu_obu, + try_read_ivf_file_header, read_ivf_frame_header, read_obus_from_ivf_frame, +}; use super::{CliOptions, IoFormat, general_read_write, input_from_either}; use general_read_write::{DoviProcessor, DoviWriter}; +fn is_av1_input(path: &Path) -> bool { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + pub struct Remover { format: IoFormat, input: PathBuf, @@ -23,9 +37,15 @@ impl Remover { } = args; let input = input_from_either("remove", input, input_pos)?; - let format = hevc_parser::io::format_from_path(&input)?; - let output = output.unwrap_or(PathBuf::from("BL.hevc")); + let (format, default_output) = if is_av1_input(&input) { + let ext = input.extension().and_then(|e| e.to_str()).unwrap_or("av1"); + (IoFormat::Raw, PathBuf::from(format!("BL_no_dovi.{ext}"))) + } else { + (hevc_parser::io::format_from_path(&input)?, PathBuf::from("BL.hevc")) + }; + + let output = output.unwrap_or(default_output); Ok(Self { format, @@ -40,6 +60,10 @@ impl Remover { } fn process_input(&self, options: CliOptions) -> Result<()> { + if is_av1_input(&self.input) { + return self.remove_from_av1(); + } + let pb = super::initialize_progress_bar(&self.format, &self.input)?; match self.format { @@ -48,6 +72,53 @@ impl Remover { } } + fn remove_from_av1(&self) -> Result<()> { + println!("Removing DoVi RPU from AV1 bitstream..."); + + let in_file = File::open(&self.input)?; + let mut reader = BufReader::new(in_file); + + if let Some(ivf_header) = try_read_ivf_file_header(&mut reader)? { + // IVF container + let out_file = BufWriter::new(File::create(&self.output).expect("Can't create file")); + let mut ivf_writer = IvfWriter::new(out_file, &ivf_header)?; + + while let Some(frame_hdr) = read_ivf_frame_header(&mut reader)? { + let mut frame_data = vec![0u8; frame_hdr.frame_size as usize]; + std::io::Read::read_exact(&mut reader, &mut frame_data)?; + + let obus = read_obus_from_ivf_frame(frame_data)?; + let mut new_frame: Vec = Vec::new(); + + for obu in &obus { + if !is_dovi_rpu_obu(obu) { + new_frame.extend_from_slice(&obu.raw_bytes); + } + } + + ivf_writer.write_frame(frame_hdr.timestamp, &new_frame)?; + } + + ivf_writer.flush()?; + } else { + // Raw AV1 bitstream + let out_file = BufWriter::new(File::create(&self.output).expect("Can't create file")); + let mut obu_writer = ObuWriter::new(out_file); + let mut obu_reader = ObuReader::new(reader); + + while let Some(obu) = obu_reader.next_obu()? { + if !is_dovi_rpu_obu(&obu) { + obu_writer.write_raw(&obu.raw_bytes)?; + } + } + + obu_writer.flush()?; + } + + println!("Done."); + Ok(()) + } + fn remove_from_raw_hevc(&self, pb: ProgressBar, options: CliOptions) -> Result<()> { let bl_out = Some(self.output.as_path()); diff --git a/src/dovi/rpu_extractor.rs b/src/dovi/rpu_extractor.rs index cc00892..10cf5c9 100644 --- a/src/dovi/rpu_extractor.rs +++ b/src/dovi/rpu_extractor.rs @@ -1,6 +1,8 @@ use anyhow::Result; use indicatif::ProgressBar; -use std::path::PathBuf; +use std::fs::File; +use std::io::{BufReader, BufWriter, Write}; +use std::path::{Path, PathBuf}; use crate::commands::ExtractRpuArgs; @@ -11,6 +13,15 @@ use super::{ }; use general_read_write::{DoviProcessor, DoviWriter}; +use super::av1::{ + OBU_METADATA, ObuReader, + is_dovi_rpu_obu, extract_dovi_t35_payload, + try_read_ivf_file_header, read_ivf_frame_header, read_obus_from_ivf_frame, +}; +use dolby_vision::rpu::dovi_rpu::DoviRpu; +use hevc_parser::hevc::{NAL_UNSPEC62, NALUnit}; +use hevc_parser::io::StartCodePreset; + pub struct RpuExtractor { format: IoFormat, input: PathBuf, @@ -18,6 +29,13 @@ pub struct RpuExtractor { limit: Option, } +fn is_av1_input(path: &Path) -> bool { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + impl RpuExtractor { pub fn from_args(args: ExtractRpuArgs) -> Result { let ExtractRpuArgs { @@ -28,7 +46,13 @@ impl RpuExtractor { } = args; let input = input_from_either("extract-rpu", input, input_pos)?; - let format = hevc_parser::io::format_from_path(&input)?; + + // For AV1 inputs use a dummy format; for HEVC use the existing detection + let format = if is_av1_input(&input) { + IoFormat::Raw + } else { + hevc_parser::io::format_from_path(&input)? + }; let rpu_out = match rpu_out { Some(path) => path, @@ -49,8 +73,93 @@ impl RpuExtractor { } fn process_input(&self, options: CliOptions) -> Result<()> { - let pb = super::initialize_progress_bar(&self.format, &self.input)?; - self.extract_rpu_from_el(pb, options) + if is_av1_input(&self.input) { + self.extract_rpu_from_av1() + } else { + let pb = super::initialize_progress_bar(&self.format, &self.input)?; + self.extract_rpu_from_el(pb, options) + } + } + + fn extract_rpu_from_av1(&self) -> Result<()> { + println!("Extracting RPU from AV1 bitstream..."); + + let file = File::open(&self.input)?; + let mut reader = BufReader::new(file); + + let mut rpus: Vec> = Vec::new(); + let mut frame_count: u64 = 0; + + // Detect IVF container by peeking at first bytes + if let Some(_ivf_header) = try_read_ivf_file_header(&mut reader)? { + // IVF container: iterate over IVF frames + while let Some(frame_hdr) = read_ivf_frame_header(&mut reader)? { + if let Some(limit) = self.limit { + if frame_count >= limit { + break; + } + } + + let mut frame_data = vec![0u8; frame_hdr.frame_size as usize]; + std::io::Read::read_exact(&mut reader, &mut frame_data)?; + + let obus = read_obus_from_ivf_frame(frame_data)?; + for obu in &obus { + if is_dovi_rpu_obu(obu) { + if let Some(t35_payload) = extract_dovi_t35_payload(&obu.payload) { + let rpu = DoviRpu::parse_itu_t35_dovi_metadata_obu(t35_payload)?; + rpus.push(rpu.write_hevc_unspec62_nalu()?); + } + } + } + + frame_count += 1; + } + } else { + // Raw AV1 bitstream + let mut obu_reader = ObuReader::new(reader); + while let Some(obu) = obu_reader.next_obu()? { + if let Some(limit) = self.limit { + if frame_count >= limit { + break; + } + } + + if obu.obu_type == OBU_METADATA && is_dovi_rpu_obu(&obu) { + if let Some(t35_payload) = extract_dovi_t35_payload(&obu.payload) { + let rpu = DoviRpu::parse_itu_t35_dovi_metadata_obu(t35_payload)?; + rpus.push(rpu.write_hevc_unspec62_nalu()?); + } + frame_count += 1; + } + } + } + + println!("Found {} RPU(s).", rpus.len()); + self.write_av1_rpu_file(&rpus) + } + + fn write_av1_rpu_file(&self, rpus: &[Vec]) -> Result<()> { + println!("Writing RPU file..."); + let mut writer = BufWriter::with_capacity( + 100_000, + File::create(&self.rpu_out).expect("Can't create file"), + ); + + for encoded_rpu in rpus { + // encoded_rpu is write_hevc_unspec62_nalu() output: starts with 0x7C 0x01 + // Same format as HEVC path: [00 00 00 01] + rpu[2..] + NALUnit::write_with_preset( + &mut writer, + &encoded_rpu[2..], + StartCodePreset::Four, + NAL_UNSPEC62, + true, + )?; + } + + writer.flush()?; + Ok(()) } fn extract_rpu_from_el(&self, pb: ProgressBar, options: CliOptions) -> Result<()> { diff --git a/src/dovi/rpu_injector.rs b/src/dovi/rpu_injector.rs index 733bbf0..aa3b149 100644 --- a/src/dovi/rpu_injector.rs +++ b/src/dovi/rpu_injector.rs @@ -1,6 +1,6 @@ use std::fs::File; -use std::io::{BufReader, BufWriter, Write, stdout}; -use std::path::PathBuf; +use std::io::{BufReader, BufWriter, Read, Write, stdout}; +use std::path::{Path, PathBuf}; use anyhow::{Result, bail}; use indicatif::ProgressBar; @@ -14,9 +14,234 @@ use dolby_vision::rpu::utils::parse_rpu_file; use crate::commands::InjectRpuArgs; +use super::av1::{ + IvfFrameHeader, IvfWriter, Obu, OBU_TEMPORAL_DELIMITER, + build_dovi_obu, is_dovi_rpu_obu, + try_read_ivf_file_header, read_ivf_frame_header, read_obus_from_ivf_frame, +}; use super::hdr10plus_utils::prefix_sei_removed_hdr10plus_nalu; use super::{CliOptions, DoviRpu, IoFormat, input_from_either}; +fn is_av1_input(path: &Path) -> bool { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + +fn inject_rpu_av1(input: &Path, rpu_in: &Path, output: &Path) -> Result<()> { + println!("Parsing RPU file..."); + stdout().flush().ok(); + + let rpus = parse_rpu_file(rpu_in)?; + println!("Loaded {} RPU(s).", rpus.len()); + + let file = File::open(input)?; + let mut reader = BufReader::with_capacity(100_000, file); + + if let Some(ivf_header) = try_read_ivf_file_header(&mut reader)? { + // IVF container — IvfWriter owns frame-level I/O (consistent with remover) + let out_file = BufWriter::with_capacity( + 100_000, + File::create(output).expect("Can't create output file"), + ); + let mut ivf_writer = IvfWriter::new(out_file, &ivf_header)?; + inject_ivf_av1(&mut reader, &mut ivf_writer, &rpus)?; + ivf_writer.flush()?; + } else { + // Raw AV1 bitstream + let mut writer = BufWriter::with_capacity( + 100_000, + File::create(output).expect("Can't create output file"), + ); + inject_raw_av1(&mut reader, &mut writer, &rpus)?; + writer.flush()?; + } + + println!("Rewriting with interleaved RPU OBUs: Done."); + Ok(()) +} + +fn inject_ivf_av1( + reader: &mut R, + ivf_writer: &mut IvfWriter, + rpus: &[DoviRpu], +) -> Result<()> { + let total_rpus = rpus.len(); + let mut tu_index = 0usize; + let mut warned_existing = false; + let mut warned_mismatch = false; + + loop { + let fh: IvfFrameHeader = match read_ivf_frame_header(reader)? { + Some(h) => h, + None => break, + }; + + let mut frame_data = vec![0u8; fh.frame_size as usize]; + reader.read_exact(&mut frame_data)?; + + let obus = read_obus_from_ivf_frame(frame_data)?; + + if !warned_existing && obus.iter().any(|o| is_dovi_rpu_obu(o)) { + warned_existing = true; + println!( + "\nWarning: Input file already has Dolby Vision RPU OBUs; \ + they will be replaced." + ); + } + + let encoded = if tu_index < total_rpus { + build_dovi_obu(&rpus[tu_index])? + } else { + if !warned_mismatch { + warned_mismatch = true; + println!( + "\nWarning: mismatched lengths. \ + RPU has {total_rpus} entries but video has more frames. \ + Last RPU will be duplicated." + ); + } + match rpus.last() { + Some(rpu) => build_dovi_obu(rpu)?, + None => bail!("No RPU available for TU {tu_index}"), + } + }; + + let output_frame = build_output_frame_av1(&obus, &encoded); + ivf_writer.write_frame(fh.timestamp, &output_frame)?; + + tu_index += 1; + } + + if tu_index < total_rpus { + println!( + "\nWarning: mismatched lengths. RPU has {total_rpus} entries \ + but video has {tu_index} frames. Excess RPU data was ignored." + ); + } + + Ok(()) +} + +fn inject_raw_av1( + reader: &mut R, + writer: &mut W, + rpus: &[DoviRpu], +) -> Result<()> { + let total_rpus = rpus.len(); + let mut tu_index = 0usize; + let mut warned_existing = false; + let mut warned_mismatch = false; + + let mut current_td: Option = None; + let mut pending: Vec = Vec::new(); + + loop { + let obu_opt = Obu::read_from(reader)?; + let is_eof = obu_opt.is_none(); + let is_td = obu_opt + .as_ref() + .map(|o| o.obu_type == OBU_TEMPORAL_DELIMITER) + .unwrap_or(false); + + if (is_eof || is_td) && current_td.is_some() { + if !warned_existing && pending.iter().any(|o| is_dovi_rpu_obu(o)) { + warned_existing = true; + println!( + "\nWarning: Input file already has Dolby Vision RPU OBUs; \ + they will be replaced." + ); + } + + let encoded = if tu_index < total_rpus { + build_dovi_obu(&rpus[tu_index])? + } else { + if !warned_mismatch { + warned_mismatch = true; + println!( + "\nWarning: mismatched lengths. \ + RPU has {total_rpus} entries but video has more frames. \ + Last RPU will be duplicated." + ); + } + match rpus.last() { + Some(rpu) => build_dovi_obu(rpu)?, + None => bail!("No RPU available for TU {tu_index}"), + } + }; + + // Write: TD + RPU OBU + remaining OBUs (skip existing DoVi) + let td = current_td.take().unwrap(); + writer.write_all(&td.raw_bytes)?; + writer.write_all(&encoded)?; + for obu in pending.drain(..) { + if !is_dovi_rpu_obu(&obu) { + writer.write_all(&obu.raw_bytes)?; + } + } + + tu_index += 1; + } + + match obu_opt { + None => break, + Some(obu) => { + if obu.obu_type == OBU_TEMPORAL_DELIMITER { + current_td = Some(obu); + pending.clear(); + } else if current_td.is_some() { + pending.push(obu); + } else { + // OBUs before the first TD — pass through unchanged + writer.write_all(&obu.raw_bytes)?; + } + } + } + } + + if tu_index < total_rpus { + println!( + "\nWarning: mismatched lengths. RPU has {total_rpus} entries \ + but video has {tu_index} frames. Excess RPU data was ignored." + ); + } + + Ok(()) +} + +/// Build the output byte buffer for one IVF temporal unit: +/// inject the RPU OBU right after OBU_TEMPORAL_DELIMITER (if present) +/// and strip any existing Dolby Vision RPU OBUs. +fn build_output_frame_av1(obus: &[Obu], encoded: &[u8]) -> Vec { + let mut out = Vec::new(); + let mut injected = false; + + // Insertion point: right after OBU_TEMPORAL_DELIMITER, or at position 0 + let insert_after_td = obus + .iter() + .position(|o| o.obu_type == OBU_TEMPORAL_DELIMITER) + .map(|i| i + 1) + .unwrap_or(0); + + for (i, obu) in obus.iter().enumerate() { + if !injected && i == insert_after_td { + out.extend_from_slice(encoded); + injected = true; + } + if is_dovi_rpu_obu(obu) { + continue; // drop existing Dolby Vision RPU + } + out.extend_from_slice(&obu.raw_bytes); + } + + if !injected { + out.extend_from_slice(encoded); + } + + out +} + pub struct RpuInjector { input: PathBuf, rpu_in: PathBuf, @@ -91,6 +316,15 @@ impl RpuInjector { pub fn inject_rpu(args: InjectRpuArgs, cli_options: CliOptions) -> Result<()> { let input = input_from_either("inject-rpu", args.input.clone(), args.input_pos.clone())?; + + if is_av1_input(&input) { + let output = args.output.clone().unwrap_or_else(|| { + let ext = input.extension().and_then(|e| e.to_str()).unwrap_or("av1"); + PathBuf::from(format!("injected_output.{ext}")) + }); + return inject_rpu_av1(&input, &args.rpu_in, &output); + } + let format = hevc_parser::io::format_from_path(&input)?; if let IoFormat::Raw = format {