From 66356593e31c8285ce57c1ba7c84ba8d7e6079dc Mon Sep 17 00:00:00 2001 From: sven-pke Date: Thu, 19 Mar 2026 12:31:51 +0100 Subject: [PATCH 1/2] Add AV1 bitstream support alongside existing HEVC Introduces an 'av1_parser' workspace crate that owns all generic AV1 OBU parsing and I/O (ObuReader, IvfWriter, ObuWriter, LEB128, IVF container handling). All existing HEVC functionality is preserved. Commands that process bitstreams (extract, inject, remove) now auto-detect the input format: .av1 / .ivf -> new AV1 code path everything else -> existing HEVC code path (unchanged) The hdr10plus library crate gains a new 'av1' module with HDR10+ OBU encoding/decoding (OBU_METADATA T.35, provider_code=0x003C). HDR10+-specific AV1 helpers live in src/core/av1_parser.rs. --- Cargo.lock | 8 ++ Cargo.toml | 4 + av1_parser/Cargo.toml | 7 + av1_parser/src/lib.rs | 299 +++++++++++++++++++++++++++++++++++++++ hdr10plus/src/av1/mod.rs | 92 ++++++++++++ hdr10plus/src/lib.rs | 2 + src/commands/extract.rs | 168 +++++++++++++++++++--- src/commands/inject.rs | 256 +++++++++++++++++++++++++++++++-- src/commands/remove.rs | 120 +++++++++++++--- src/core/av1_parser.rs | 249 ++++++++++++++++++++++++++++++++ src/core/mod.rs | 1 + 11 files changed, 1163 insertions(+), 43 deletions(-) create mode 100644 av1_parser/Cargo.toml create mode 100644 av1_parser/src/lib.rs create mode 100644 hdr10plus/src/av1/mod.rs create mode 100644 src/core/av1_parser.rs diff --git a/Cargo.lock b/Cargo.lock index d10a879..6072119 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -125,6 +125,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" @@ -618,6 +625,7 @@ dependencies = [ "anyhow", "assert_cmd", "assert_fs", + "av1_parser", "bitvec_helpers", "clap", "hdr10plus", diff --git a/Cargo.toml b/Cargo.toml index 476c0b9..f374233 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ license = "MIT" [dependencies] hdr10plus = { path = "./hdr10plus", features = ["hevc", "json"] } +av1_parser = { path = "av1_parser" } bitvec_helpers = { version = "4.0.1", default-features = false, features = ["bitstream-io"] } hevc_parser = { version = "0.6.10", features = ["hevc_io"] } @@ -30,8 +31,11 @@ path = "src/main.rs" [workspace] members = [ + ".", "hdr10plus", + "av1_parser", ] +resolver = "2" [features] default = ["system-font"] 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/hdr10plus/src/av1/mod.rs b/hdr10plus/src/av1/mod.rs new file mode 100644 index 0000000..f2d42c2 --- /dev/null +++ b/hdr10plus/src/av1/mod.rs @@ -0,0 +1,92 @@ +use anyhow::Result; + +use crate::metadata::{Hdr10PlusMetadata, Hdr10PlusMetadataEncOpts}; + +#[cfg(feature = "json")] +use crate::metadata_json::Hdr10PlusJsonMetadata; + +/// OBU type for metadata +const OBU_METADATA: u8 = 5; + +/// T.35 metadata type value for ITU-T T.35 +const METADATA_TYPE_ITUT_T35: u64 = 4; + +/// Encodes HDR10+ metadata as a complete AV1 OBU_METADATA unit. +/// +/// The OBU contains: +/// - OBU header (1 byte, type=5, has_size_field=1) +/// - OBU size (LEB128) +/// - metadata_type = 4 (METADATA_TYPE_ITUT_T35, LEB128) +/// - ITU-T T.35 payload (country code + HDR10+ bitstream) +pub fn encode_hdr10plus_obu(metadata: &Hdr10PlusMetadata, validate: bool) -> Result> { + let opts = Hdr10PlusMetadataEncOpts { + validate, + with_country_code: true, + ..Default::default() + }; + + // T.35 payload including country code (0xB5), terminal_provider_code, etc. + let t35_payload = metadata.encode_with_opts(&opts)?; + + // OBU_METADATA payload: metadata_type (LEB128) + T.35 payload + let mut obu_payload = encode_leb128(METADATA_TYPE_ITUT_T35); + obu_payload.extend_from_slice(&t35_payload); + + // 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) +} + +#[cfg(feature = "json")] +pub fn encode_av1_from_json(metadata: &Hdr10PlusJsonMetadata, validate: bool) -> Result> { + let meta = Hdr10PlusMetadata::try_from(metadata)?; + encode_hdr10plus_obu(&meta, validate) +} + +/// 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) +} diff --git a/hdr10plus/src/lib.rs b/hdr10plus/src/lib.rs index c56dad6..455c849 100644 --- a/hdr10plus/src/lib.rs +++ b/hdr10plus/src/lib.rs @@ -6,6 +6,8 @@ pub mod metadata_json; #[cfg(feature = "hevc")] pub mod hevc; +pub mod av1; + /// C API module #[cfg(any(cargo_c, feature = "capi"))] pub mod capi; diff --git a/src/commands/extract.rs b/src/commands/extract.rs index 83bb00f..399f5e1 100644 --- a/src/commands/extract.rs +++ b/src/commands/extract.rs @@ -1,11 +1,30 @@ -use anyhow::Result; +use std::fs::File; +use std::io::{BufReader, BufWriter, Read, Write, stdout}; +use std::path::Path; + +use anyhow::{Result, bail}; + +use hdr10plus::metadata::Hdr10PlusMetadata; +use hdr10plus::metadata_json::generate_json; use super::{CliOptions, ExtractArgs, input_from_either}; +use crate::core::ParserError; +use crate::core::av1_parser::{ + Av1NaluParser, Obu, OBU_METADATA, extract_hdr10plus_t35_bytes, + read_ivf_frame_header, read_obus_from_ivf_frame, try_read_ivf_file_header, +}; use crate::core::initialize_progress_bar; -use crate::core::parser::{Parser, ParserOptions}; +use crate::core::parser::{Parser, ParserOptions, TOOL_NAME, TOOL_VERSION}; pub struct Extractor {} +fn is_av1_input(path: &Path) -> bool { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + impl Extractor { pub fn extract_json(args: ExtractArgs, mut options: CliOptions) -> Result<()> { let ExtractArgs { @@ -17,22 +36,139 @@ impl Extractor { } = args; let input = input_from_either("extract", input, input_pos)?; - let format = hevc_parser::io::format_from_path(&input)?; + if is_av1_input(&input) { + Self::extract_json_av1(input, output, options, limit) + } else { + let format = hevc_parser::io::format_from_path(&input)?; - if !options.verify && output.is_none() { - options.verify = true - }; + if !options.verify && output.is_none() { + options.verify = true + }; - let pb = initialize_progress_bar(&format, &input)?; - let mut parser = Parser::new( - input, - output, - options, - pb, - skip_reorder, - ParserOptions { limit }, - ); + let pb = initialize_progress_bar(&format, &input)?; + let mut parser = Parser::new( + input, + output, + options, + pb, + skip_reorder, + ParserOptions { limit }, + ); + + parser.process_input(&format) + } + } + + fn extract_json_av1( + input: std::path::PathBuf, + output: Option, + options: CliOptions, + limit: Option, + ) -> Result<()> { + let file = File::open(&input)?; + let mut reader = BufReader::with_capacity(100_000, file); + + let mut av1_nalu_parser = Av1NaluParser::new(); + let mut t35_frames: Vec> = Vec::new(); + let mut obu_count = 0u64; + + let is_ivf = try_read_ivf_file_header(&mut reader)?.is_some(); + + if is_ivf { + loop { + let fh = match read_ivf_frame_header(&mut 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)?; + for obu in &obus { + av1_nalu_parser.process_obu(obu)?; + if obu.obu_type == OBU_METADATA { + if let Some(t35) = + extract_hdr10plus_t35_bytes(&obu.payload, options.validate) + { + t35_frames.push(t35); + } + } + obu_count += 1; + if limit.map(|l| obu_count >= l).unwrap_or(false) { + break; + } + } + + if limit.map(|l| obu_count >= l).unwrap_or(false) { + break; + } + } + } else { + loop { + match Obu::read_from(&mut reader) { + Ok(Some(obu)) => { + av1_nalu_parser.process_obu(&obu)?; + + if obu.obu_type == OBU_METADATA { + if let Some(t35) = + extract_hdr10plus_t35_bytes(&obu.payload, options.validate) + { + t35_frames.push(t35); + } + } + + obu_count += 1; + if limit.map(|l| obu_count >= l).unwrap_or(false) { + break; + } + } + Ok(None) => break, + Err(e) => return Err(e), + } + } + } + + if t35_frames.is_empty() { + bail!(ParserError::NoMetadataFound); + } + + if options.verify { + bail!(ParserError::MetadataDetected); + } + + print!("Reading parsed dynamic metadata... "); + stdout().flush().ok(); + + let mut complete_metadata: Vec = Vec::new(); + for t35_bytes in &t35_frames { + let meta = Hdr10PlusMetadata::parse(t35_bytes)?; + if options.validate { + meta.validate()?; + } + complete_metadata.push(meta); + } + + println!("Done."); + + match output { + Some(path) => { + let save_file = File::create(&path).expect("Can't create file"); + let mut writer = BufWriter::with_capacity(10_000_000, save_file); + + print!("Generating and writing metadata to JSON file... "); + stdout().flush().ok(); + + let list: Vec<&Hdr10PlusMetadata> = complete_metadata.iter().collect(); + let final_json = generate_json(&list, TOOL_NAME, TOOL_VERSION); + + writeln!(writer, "{}", serde_json::to_string_pretty(&final_json)?)?; + writer.flush()?; + + println!("Done."); + } + None => bail!("Output path required for AV1 extraction"), + } - parser.process_input(&format) + Ok(()) } } diff --git a/src/commands/inject.rs b/src/commands/inject.rs index 07b40b4..8873eb4 100644 --- a/src/commands/inject.rs +++ b/src/commands/inject.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 hevc_parser::utils::{ @@ -19,6 +19,13 @@ use crate::core::{initialize_progress_bar, st2094_40_sei_msg}; use super::{CliOptions, input_from_either}; +fn is_av1_input(path: &Path) -> bool { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + pub struct Injector { input: PathBuf, json_in: PathBuf, @@ -96,16 +103,249 @@ impl Injector { pub fn inject_json(args: InjectArgs, cli_options: CliOptions) -> Result<()> { let input = input_from_either("inject", args.input.clone(), args.input_pos.clone())?; - let format = hevc_parser::io::format_from_path(&input)?; - if let IoFormat::Raw = format { - let mut injector = Injector::from_args(args, cli_options)?; + if is_av1_input(&input) { + Self::inject_json_av1(args, cli_options) + } else { + let format = hevc_parser::io::format_from_path(&input)?; + + if let IoFormat::Raw = format { + let mut injector = Injector::from_args(args, cli_options)?; + + injector.process_input()?; + injector.interleave_hdr10plus_nals() + } else { + bail!("Injector: Must be a raw HEVC bitstream file") + } + } + } + + fn inject_json_av1(args: InjectArgs, cli_options: CliOptions) -> Result<()> { + use crate::core::av1_parser::{ + Av1NaluParser, IvfFrameHeader, Obu, OBU_TEMPORAL_DELIMITER, is_hdr10plus_obu, + read_ivf_frame_header, read_obus_from_ivf_frame, try_read_ivf_file_header, + write_ivf_frame_header, + }; + use hdr10plus::av1::encode_av1_from_json; + + let InjectArgs { + input, + input_pos, + json, + output, + } = args; + let input = input_from_either("inject", input, input_pos)?; + let output = output.unwrap_or_else(|| PathBuf::from("injected_output.av1")); + + println!("Parsing JSON file..."); + stdout().flush().ok(); + + let metadata_root = MetadataJsonRoot::from_file(&json)?; + let metadata_list: Vec = metadata_root.scene_info; + + if metadata_list.is_empty() { + bail!("Empty HDR10+ SceneInfo array"); + } + + let file = File::open(&input)?; + let mut reader = BufReader::with_capacity(100_000, file); + + let out_file = File::create(&output).expect("Can't create output file"); + let mut writer = BufWriter::with_capacity(100_000, out_file); + + let total_meta = metadata_list.len(); + + if let Some(ivf_header) = try_read_ivf_file_header(&mut reader)? { + writer.write_all(&ivf_header)?; + + let mut tu_index = 0usize; + let mut last_encoded: Option> = None; + let mut warned_existing = false; + let mut warned_mismatch = false; + + loop { + let fh: IvfFrameHeader = match read_ivf_frame_header(&mut 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_hdr10plus_obu(o, cli_options.validate)) + { + warned_existing = true; + println!( + "\nWarning: Input file already has HDR10+ metadata OBUs; \ + they will be replaced." + ); + } + + let encoded = if tu_index < total_meta { + let enc = + encode_av1_from_json(&metadata_list[tu_index], cli_options.validate)?; + last_encoded = Some(enc.clone()); + enc + } else { + if !warned_mismatch { + warned_mismatch = true; + println!( + "\nWarning: mismatched lengths. \ + Metadata has {total_meta} entries but video has more frames. \ + Last metadata will be duplicated." + ); + } + match &last_encoded { + Some(enc) => enc.clone(), + None => bail!("No HDR10+ metadata available for TU {tu_index}"), + } + }; + + let output_frame = + Self::build_av1_output_frame(&obus, &encoded, cli_options.validate); + + write_ivf_frame_header(&mut writer, output_frame.len() as u32, fh.timestamp)?; + writer.write_all(&output_frame)?; - injector.process_input()?; - injector.interleave_hdr10plus_nals() + tu_index += 1; + } + + if tu_index < total_meta { + println!( + "\nWarning: mismatched lengths. Metadata has {total_meta} entries \ + but video has {tu_index} frames. Excess metadata was ignored." + ); + } } else { - bail!("Injector: Must be a raw HEVC bitstream file") + // Raw OBU stream + let mut av1_parser = Av1NaluParser::new(); + let mut tu_index = 0usize; + let mut last_encoded: Option> = None; + 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(&mut 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_hdr10plus_obu(o, cli_options.validate)) + { + warned_existing = true; + println!( + "\nWarning: Input file already has HDR10+ metadata OBUs; \ + they will be replaced." + ); + } + + let encoded = if tu_index < total_meta { + let enc = encode_av1_from_json( + &metadata_list[tu_index], + cli_options.validate, + )?; + last_encoded = Some(enc.clone()); + enc + } else { + if !warned_mismatch { + warned_mismatch = true; + println!( + "\nWarning: mismatched lengths. \ + Metadata has {total_meta} entries but video has more frames. \ + Last metadata will be duplicated." + ); + } + match &last_encoded { + Some(enc) => enc.clone(), + None => bail!("No HDR10+ metadata available for TU {tu_index}"), + } + }; + + let td = current_td.take().unwrap(); + writer.write_all(&td.raw_bytes)?; + writer.write_all(&encoded)?; + for obu in pending.drain(..) { + if !is_hdr10plus_obu(&obu, cli_options.validate) { + writer.write_all(&obu.raw_bytes)?; + } + } + + tu_index += 1; + } + + match obu_opt { + None => break, + Some(obu) => { + av1_parser.process_obu(&obu)?; + if obu.obu_type == OBU_TEMPORAL_DELIMITER { + current_td = Some(obu); + pending.clear(); + } else if current_td.is_some() { + pending.push(obu); + } else { + writer.write_all(&obu.raw_bytes)?; + } + } + } + } + + if tu_index < total_meta { + println!( + "\nWarning: mismatched lengths. Metadata has {total_meta} entries \ + but video has {tu_index} frames. Excess metadata was ignored." + ); + } + } + + println!("Rewriting with interleaved HDR10+ metadata OBUs: Done."); + writer.flush()?; + Ok(()) + } + + fn build_av1_output_frame( + obus: &[crate::core::av1_parser::Obu], + encoded: &[u8], + validate: bool, + ) -> Vec { + use crate::core::av1_parser::{OBU_TEMPORAL_DELIMITER, is_hdr10plus_obu}; + + let mut out = Vec::new(); + let mut injected = false; + + 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_hdr10plus_obu(obu, validate) { + continue; + } + out.extend_from_slice(&obu.raw_bytes); + } + + if !injected { + out.extend_from_slice(encoded); } + + out } fn process_input(&mut self) -> Result<()> { diff --git a/src/commands/remove.rs b/src/commands/remove.rs index 310e847..9751891 100644 --- a/src/commands/remove.rs +++ b/src/commands/remove.rs @@ -1,6 +1,6 @@ use std::fs::File; -use std::io::{BufRead, BufReader, BufWriter, Write}; -use std::path::PathBuf; +use std::io::{BufRead, BufReader, BufWriter, Read, Write}; +use std::path::{Path, PathBuf}; use anyhow::{Result, bail}; use hevc_parser::HevcParser; @@ -13,6 +13,13 @@ use hevc_parser::io::{IoFormat, IoProcessor}; use super::{CliOptions, RemoveArgs, input_from_either}; use crate::core::{initialize_progress_bar, prefix_sei_removed_hdr10plus_nalu}; +fn is_av1_input(path: &Path) -> bool { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + pub struct Remover { input: PathBuf, progress_bar: ProgressBar, @@ -20,7 +27,7 @@ pub struct Remover { } impl Remover { - pub fn remove_sei(args: RemoveArgs, _options: CliOptions) -> Result<()> { + pub fn remove_sei(args: RemoveArgs, options: CliOptions) -> Result<()> { let RemoveArgs { input, input_pos, @@ -28,29 +35,104 @@ impl Remover { } = args; let input = input_from_either("remove", input, input_pos)?; - let format = hevc_parser::io::format_from_path(&input)?; + if is_av1_input(&input) { + Self::remove_sei_av1(input, output, options) + } else { + let format = hevc_parser::io::format_from_path(&input)?; + + if format == IoFormat::Matroska { + bail!("Remover: Matroska format unsupported"); + } + + let hevc_out = match output { + Some(path) => path, + None => PathBuf::from("hdr10plus_removed_output.hevc"), + }; + + let pb = initialize_progress_bar(&format, &input)?; - if format == IoFormat::Matroska { - bail!("Remover: Matroska format unsupported"); + let mut remover = Remover { + input, + progress_bar: pb, + writer: BufWriter::with_capacity( + 100_000, + File::create(hevc_out).expect("Can't create file"), + ), + }; + + remover.process_input(&format) } + } - let hevc_out = match output { - Some(path) => path, - None => PathBuf::from("hdr10plus_removed_output.hevc"), + fn remove_sei_av1( + input: PathBuf, + output: Option, + options: CliOptions, + ) -> Result<()> { + use crate::core::av1_parser::{ + Obu, is_hdr10plus_obu, read_ivf_frame_header, read_obus_from_ivf_frame, + try_read_ivf_file_header, write_ivf_frame_header, }; - let pb = initialize_progress_bar(&format, &input)?; + let out_path = output.unwrap_or_else(|| PathBuf::from("hdr10plus_removed_output.av1")); - let mut remover = Remover { - input, - progress_bar: pb, - writer: BufWriter::with_capacity( - 100_000, - File::create(hevc_out).expect("Can't create file"), - ), - }; + let file = File::open(&input)?; + let file_len = file.metadata().map(|m| m.len()).unwrap_or(0); + let mut reader = BufReader::with_capacity(100_000, file); + + let out_file = File::create(&out_path).expect("Can't create output file"); + let mut writer = BufWriter::with_capacity(100_000, out_file); + + let pb = initialize_progress_bar(&IoFormat::Raw, &input)?; + let mut bytes_read = 0u64; - remover.process_input(&format) + if let Some(ivf_header) = try_read_ivf_file_header(&mut reader)? { + writer.write_all(&ivf_header)?; + bytes_read += ivf_header.len() as u64; + + loop { + let fh = match read_ivf_frame_header(&mut reader)? { + Some(h) => h, + None => break, + }; + bytes_read += 12; + pb.set_position(bytes_read * 100 / file_len.max(1)); + + let mut frame_data = vec![0u8; fh.frame_size as usize]; + reader.read_exact(&mut frame_data)?; + bytes_read += fh.frame_size as u64; + + let obus = read_obus_from_ivf_frame(frame_data)?; + + let output_frame: Vec = obus + .iter() + .filter(|o| !is_hdr10plus_obu(o, options.validate)) + .flat_map(|o| o.raw_bytes.iter().copied()) + .collect(); + + write_ivf_frame_header(&mut writer, output_frame.len() as u32, fh.timestamp)?; + writer.write_all(&output_frame)?; + } + } else { + loop { + match Obu::read_from(&mut reader) { + Ok(Some(obu)) => { + bytes_read += obu.raw_bytes.len() as u64; + pb.set_position(bytes_read * 100 / file_len.max(1)); + + if !is_hdr10plus_obu(&obu, options.validate) { + writer.write_all(&obu.raw_bytes)?; + } + } + Ok(None) => break, + Err(e) => return Err(e), + } + } + } + + pb.finish_and_clear(); + writer.flush()?; + Ok(()) } pub fn process_input(&mut self, format: &IoFormat) -> Result<()> { diff --git a/src/core/av1_parser.rs b/src/core/av1_parser.rs new file mode 100644 index 0000000..02e6656 --- /dev/null +++ b/src/core/av1_parser.rs @@ -0,0 +1,249 @@ +#![allow(dead_code, unused_imports)] + +use anyhow::Result; +use bitvec_helpers::bitstream_io_reader::BsIoSliceReader; + +// Re-export all generic types/functions from the crate +pub use av1_parser::{ + IvfFrameHeader, IvfWriter, Obu, ObuReader, ObuWriter, + OBU_TEMPORAL_DELIMITER, OBU_METADATA, + OBU_SEQUENCE_HEADER, OBU_FRAME_HEADER, OBU_FRAME, OBU_REDUNDANT_FRAME_HEADER, + decode_leb128, encode_leb128, + try_read_ivf_file_header, read_ivf_frame_header, write_ivf_frame_header, + read_obus_from_ivf_frame, +}; + +// --------------------------------------------------------------------------- +// HDR10+-specific constants +// --------------------------------------------------------------------------- + +// Metadata type for ITU-T T.35 (HDR10+) +pub const METADATA_TYPE_ITUT_T35: u64 = 4; + +// HDR10+ T.35 header identifiers +pub const HDR10PLUS_COUNTRY_CODE: u8 = 0xB5; +pub const HDR10PLUS_PROVIDER_CODE: u16 = 0x003C; +pub const HDR10PLUS_ORIENTED_CODE: u16 = 0x0001; +pub const HDR10PLUS_APP_ID: u8 = 4; + +// --------------------------------------------------------------------------- +// HDR10+ detection helper +// --------------------------------------------------------------------------- + +/// Returns the T.35 payload bytes (starting at country_code = 0xB5) if this +/// OBU_METADATA payload contains HDR10+ data. The returned slice is the +/// portion that `Hdr10PlusMetadata::parse()` expects. +/// +/// Layout of an OBU_METADATA payload for HDR10+: +/// ```text +/// metadata_type (LEB128) = 4 +/// country_code (u8) = 0xB5 +/// provider_code (u16 BE) = 0x003C +/// oriented_code (u16 BE) = 0x0001 +/// app_id (u8) = 4 +/// app_version (u8) = 1 +/// +/// ``` +pub fn extract_hdr10plus_t35_bytes( + obu_payload: &[u8], + validate: bool, +) -> Option> { + if obu_payload.is_empty() { + return None; + } + + // metadata_type + let (metadata_type, mt_len) = decode_leb128(obu_payload); + if metadata_type != METADATA_TYPE_ITUT_T35 { + return None; + } + + let t35 = &obu_payload[mt_len..]; + if t35.len() < 7 { + return None; + } + + let country_code = t35[0]; + if country_code != HDR10PLUS_COUNTRY_CODE { + return None; + } + + let provider_code = u16::from_be_bytes([t35[1], t35[2]]); + let oriented_code = u16::from_be_bytes([t35[3], t35[4]]); + let app_id = t35[5]; + let app_version = t35[6]; + + if provider_code != HDR10PLUS_PROVIDER_CODE + || oriented_code != HDR10PLUS_ORIENTED_CODE + || app_id != HDR10PLUS_APP_ID + { + return None; + } + + let valid_version = if validate { + app_version == 1 + } else { + app_version <= 1 + }; + + if !valid_version { + return None; + } + + // Return T.35 bytes starting at country_code (what parse() expects) + Some(t35.to_vec()) +} + +/// Returns `true` if this OBU is an OBU_METADATA carrying HDR10+ T.35 data. +pub fn is_hdr10plus_obu(obu: &Obu, validate: bool) -> bool { + obu.obu_type == OBU_METADATA + && extract_hdr10plus_t35_bytes(&obu.payload, validate).is_some() +} + +// --------------------------------------------------------------------------- +// Stateful AV1 parser +// --------------------------------------------------------------------------- + +/// Information about a single temporal unit derived from stream parsing. +#[derive(Debug, Clone)] +pub struct TemporalUnitInfo { + /// Zero-based index of this temporal unit in the stream. + pub index: usize, + /// `true` when the temporal unit produces a visible output frame. + /// This is `false` only for frames with `show_frame = 0` and + /// `showable_frame = 0`. + pub is_displayed: bool, + /// `true` when the temporal unit reuses a previously decoded frame + /// via `show_existing_frame`. + pub is_show_existing: bool, +} + +/// Stateful AV1 bitstream parser. +/// +/// Feed it each `Obu` in stream order via [`process_obu`]. After parsing the +/// entire stream, [`temporal_units`] provides per-TU display metadata needed +/// for frame-accurate HDR10+ injection. +pub struct Av1NaluParser { + /// Derived from the sequence header; affects frame header interpretation. + pub reduced_still_picture_header: bool, + + /// Running count of temporal units seen so far (incremented on each TD). + pub temporal_unit_count: usize, + + /// Per-temporal-unit display info (populated as frame headers are parsed). + pub temporal_units: Vec, + + /// A frame header was already parsed in the current TU (to skip redundant ones). + frame_header_parsed_in_tu: bool, +} + +impl Av1NaluParser { + pub fn new() -> Self { + Self { + reduced_still_picture_header: false, + temporal_unit_count: 0, + temporal_units: Vec::new(), + frame_header_parsed_in_tu: false, + } + } + + /// Process one OBU and update parser state. + pub fn process_obu(&mut self, obu: &Obu) -> Result<()> { + match obu.obu_type { + OBU_TEMPORAL_DELIMITER => { + // Marks the beginning of a new temporal unit. + self.temporal_unit_count += 1; + self.frame_header_parsed_in_tu = false; + } + OBU_SEQUENCE_HEADER => { + self.parse_sequence_header(&obu.payload)?; + } + OBU_FRAME_HEADER | OBU_FRAME => { + if !self.frame_header_parsed_in_tu { + let info = self.parse_frame_display_info(&obu.payload)?; + let tu_idx = self.temporal_unit_count.saturating_sub(1); + self.temporal_units.push(TemporalUnitInfo { + index: tu_idx, + is_displayed: info.0, + is_show_existing: info.1, + }); + self.frame_header_parsed_in_tu = true; + } + } + OBU_REDUNDANT_FRAME_HEADER => { + // Redundant copies carry the same info — skip to avoid duplicates. + } + _ => {} + } + Ok(()) + } + + /// Return all collected temporal unit infos. + pub fn temporal_units(&self) -> &[TemporalUnitInfo] { + &self.temporal_units + } + + /// Returns the number of temporal units that produce a displayed frame. + pub fn display_frame_count(&self) -> usize { + self.temporal_units.iter().filter(|t| t.is_displayed).count() + } + + // ----------------------------------------------------------------------- + // Sequence header parsing + // ----------------------------------------------------------------------- + + fn parse_sequence_header(&mut self, payload: &[u8]) -> Result<()> { + if payload.len() < 1 { + return Ok(()); + } + let mut r = BsIoSliceReader::from_slice(payload); + + // seq_profile (3 bits) + let _seq_profile = r.read::<3, u8>()?; + // still_picture (1 bit) + let _still_picture = r.read::<1, u8>()?; + // reduced_still_picture_header (1 bit) + self.reduced_still_picture_header = r.read::<1, u8>()? != 0; + + Ok(()) + } + + // ----------------------------------------------------------------------- + // Frame header parsing (minimal — only what is needed for display flags) + // ----------------------------------------------------------------------- + + /// Returns `(is_displayed, is_show_existing)`. + fn parse_frame_display_info(&self, payload: &[u8]) -> Result<(bool, bool)> { + if self.reduced_still_picture_header { + // Section 5.9.2: reduced_still_picture_header implies + // show_existing_frame = 0, show_frame = 1. + return Ok((true, false)); + } + + if payload.is_empty() { + return Ok((true, false)); + } + + let mut r = BsIoSliceReader::from_slice(payload); + + // show_existing_frame (f(1)) + let show_existing_frame = r.read::<1, u8>()? != 0; + if show_existing_frame { + return Ok((true, true)); + } + + // frame_type (f(2)) + let _frame_type = r.read::<2, u8>()?; + + // show_frame (f(1)) + let show_frame = r.read::<1, u8>()? != 0; + + if show_frame { + Ok((true, false)) + } else { + // showable_frame (f(1)) + let showable_frame = r.read::<1, u8>()? != 0; + Ok((showable_frame, false)) + } + } +} diff --git a/src/core/mod.rs b/src/core/mod.rs index b0aee8d..3134c18 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -11,6 +11,7 @@ use hevc_parser::utils::{ add_start_code_emulation_prevention_3_byte, clear_start_code_emulation_prevention_3_byte, }; +pub mod av1_parser; pub mod parser; #[derive(Error, Debug)] From 191d5dcc829fb66ebdfb3cc860d1fdd94eb293db Mon Sep 17 00:00:00 2001 From: sven-pke Date: Thu, 19 Mar 2026 13:15:35 +0100 Subject: [PATCH 2/2] Remove dead Av1NaluParser code and clean up AV1 imports - Remove Av1NaluParser and TemporalUnitInfo structs from av1_parser.rs (they were unused dead code suppressed by #![allow(dead_code)]) - Remove unused re-exports (IvfWriter, ObuWriter) from av1_parser.rs - Move inner use blocks to file-level imports in inject.rs and remove.rs - Remove Av1NaluParser usage from extract.rs and inject.rs - Builds with zero warnings Co-Authored-By: Claude Sonnet 4.6 --- src/commands/extract.rs | 6 +- src/commands/inject.rs | 23 ++---- src/commands/remove.rs | 9 +-- src/core/av1_parser.rs | 160 +--------------------------------------- 4 files changed, 15 insertions(+), 183 deletions(-) diff --git a/src/commands/extract.rs b/src/commands/extract.rs index 399f5e1..ffe0324 100644 --- a/src/commands/extract.rs +++ b/src/commands/extract.rs @@ -10,7 +10,7 @@ use hdr10plus::metadata_json::generate_json; use super::{CliOptions, ExtractArgs, input_from_either}; use crate::core::ParserError; use crate::core::av1_parser::{ - Av1NaluParser, Obu, OBU_METADATA, extract_hdr10plus_t35_bytes, + Obu, OBU_METADATA, extract_hdr10plus_t35_bytes, read_ivf_frame_header, read_obus_from_ivf_frame, try_read_ivf_file_header, }; use crate::core::initialize_progress_bar; @@ -68,7 +68,6 @@ impl Extractor { let file = File::open(&input)?; let mut reader = BufReader::with_capacity(100_000, file); - let mut av1_nalu_parser = Av1NaluParser::new(); let mut t35_frames: Vec> = Vec::new(); let mut obu_count = 0u64; @@ -85,7 +84,6 @@ impl Extractor { let obus = read_obus_from_ivf_frame(frame_data)?; for obu in &obus { - av1_nalu_parser.process_obu(obu)?; if obu.obu_type == OBU_METADATA { if let Some(t35) = extract_hdr10plus_t35_bytes(&obu.payload, options.validate) @@ -107,8 +105,6 @@ impl Extractor { loop { match Obu::read_from(&mut reader) { Ok(Some(obu)) => { - av1_nalu_parser.process_obu(&obu)?; - if obu.obu_type == OBU_METADATA { if let Some(t35) = extract_hdr10plus_t35_bytes(&obu.payload, options.validate) diff --git a/src/commands/inject.rs b/src/commands/inject.rs index 8873eb4..6a9e6e8 100644 --- a/src/commands/inject.rs +++ b/src/commands/inject.rs @@ -12,9 +12,15 @@ use hevc_parser::io::{FrameBuffer, IoFormat, IoProcessor, NalBuffer, processor}; use hevc_parser::{HevcParser, NALUStartCode, hevc::*}; use processor::{HevcProcessor, HevcProcessorOpts}; +use hdr10plus::av1::encode_av1_from_json; use hdr10plus::metadata_json::{Hdr10PlusJsonMetadata, MetadataJsonRoot}; use crate::commands::InjectArgs; +use crate::core::av1_parser::{ + IvfFrameHeader, Obu, OBU_TEMPORAL_DELIMITER, is_hdr10plus_obu, + read_ivf_frame_header, read_obus_from_ivf_frame, try_read_ivf_file_header, + write_ivf_frame_header, +}; use crate::core::{initialize_progress_bar, st2094_40_sei_msg}; use super::{CliOptions, input_from_either}; @@ -121,13 +127,6 @@ impl Injector { } fn inject_json_av1(args: InjectArgs, cli_options: CliOptions) -> Result<()> { - use crate::core::av1_parser::{ - Av1NaluParser, IvfFrameHeader, Obu, OBU_TEMPORAL_DELIMITER, is_hdr10plus_obu, - read_ivf_frame_header, read_obus_from_ivf_frame, try_read_ivf_file_header, - write_ivf_frame_header, - }; - use hdr10plus::av1::encode_av1_from_json; - let InjectArgs { input, input_pos, @@ -221,7 +220,6 @@ impl Injector { } } else { // Raw OBU stream - let mut av1_parser = Av1NaluParser::new(); let mut tu_index = 0usize; let mut last_encoded: Option> = None; let mut warned_existing = false; @@ -288,7 +286,6 @@ impl Injector { match obu_opt { None => break, Some(obu) => { - av1_parser.process_obu(&obu)?; if obu.obu_type == OBU_TEMPORAL_DELIMITER { current_td = Some(obu); pending.clear(); @@ -314,13 +311,7 @@ impl Injector { Ok(()) } - fn build_av1_output_frame( - obus: &[crate::core::av1_parser::Obu], - encoded: &[u8], - validate: bool, - ) -> Vec { - use crate::core::av1_parser::{OBU_TEMPORAL_DELIMITER, is_hdr10plus_obu}; - + fn build_av1_output_frame(obus: &[Obu], encoded: &[u8], validate: bool) -> Vec { let mut out = Vec::new(); let mut injected = false; diff --git a/src/commands/remove.rs b/src/commands/remove.rs index 9751891..176707f 100644 --- a/src/commands/remove.rs +++ b/src/commands/remove.rs @@ -11,6 +11,10 @@ use indicatif::ProgressBar; use hevc_parser::io::{IoFormat, IoProcessor}; use super::{CliOptions, RemoveArgs, input_from_either}; +use crate::core::av1_parser::{ + Obu, is_hdr10plus_obu, read_ivf_frame_header, read_obus_from_ivf_frame, + try_read_ivf_file_header, write_ivf_frame_header, +}; use crate::core::{initialize_progress_bar, prefix_sei_removed_hdr10plus_nalu}; fn is_av1_input(path: &Path) -> bool { @@ -69,11 +73,6 @@ impl Remover { output: Option, options: CliOptions, ) -> Result<()> { - use crate::core::av1_parser::{ - Obu, is_hdr10plus_obu, read_ivf_frame_header, read_obus_from_ivf_frame, - try_read_ivf_file_header, write_ivf_frame_header, - }; - let out_path = output.unwrap_or_else(|| PathBuf::from("hdr10plus_removed_output.av1")); let file = File::open(&input)?; diff --git a/src/core/av1_parser.rs b/src/core/av1_parser.rs index 02e6656..bc0e0d1 100644 --- a/src/core/av1_parser.rs +++ b/src/core/av1_parser.rs @@ -1,14 +1,8 @@ -#![allow(dead_code, unused_imports)] - -use anyhow::Result; -use bitvec_helpers::bitstream_io_reader::BsIoSliceReader; - -// Re-export all generic types/functions from the crate +// Re-export generic types/functions from the crate pub use av1_parser::{ - IvfFrameHeader, IvfWriter, Obu, ObuReader, ObuWriter, + IvfFrameHeader, Obu, OBU_TEMPORAL_DELIMITER, OBU_METADATA, - OBU_SEQUENCE_HEADER, OBU_FRAME_HEADER, OBU_FRAME, OBU_REDUNDANT_FRAME_HEADER, - decode_leb128, encode_leb128, + decode_leb128, try_read_ivf_file_header, read_ivf_frame_header, write_ivf_frame_header, read_obus_from_ivf_frame, }; @@ -99,151 +93,3 @@ pub fn is_hdr10plus_obu(obu: &Obu, validate: bool) -> bool { obu.obu_type == OBU_METADATA && extract_hdr10plus_t35_bytes(&obu.payload, validate).is_some() } - -// --------------------------------------------------------------------------- -// Stateful AV1 parser -// --------------------------------------------------------------------------- - -/// Information about a single temporal unit derived from stream parsing. -#[derive(Debug, Clone)] -pub struct TemporalUnitInfo { - /// Zero-based index of this temporal unit in the stream. - pub index: usize, - /// `true` when the temporal unit produces a visible output frame. - /// This is `false` only for frames with `show_frame = 0` and - /// `showable_frame = 0`. - pub is_displayed: bool, - /// `true` when the temporal unit reuses a previously decoded frame - /// via `show_existing_frame`. - pub is_show_existing: bool, -} - -/// Stateful AV1 bitstream parser. -/// -/// Feed it each `Obu` in stream order via [`process_obu`]. After parsing the -/// entire stream, [`temporal_units`] provides per-TU display metadata needed -/// for frame-accurate HDR10+ injection. -pub struct Av1NaluParser { - /// Derived from the sequence header; affects frame header interpretation. - pub reduced_still_picture_header: bool, - - /// Running count of temporal units seen so far (incremented on each TD). - pub temporal_unit_count: usize, - - /// Per-temporal-unit display info (populated as frame headers are parsed). - pub temporal_units: Vec, - - /// A frame header was already parsed in the current TU (to skip redundant ones). - frame_header_parsed_in_tu: bool, -} - -impl Av1NaluParser { - pub fn new() -> Self { - Self { - reduced_still_picture_header: false, - temporal_unit_count: 0, - temporal_units: Vec::new(), - frame_header_parsed_in_tu: false, - } - } - - /// Process one OBU and update parser state. - pub fn process_obu(&mut self, obu: &Obu) -> Result<()> { - match obu.obu_type { - OBU_TEMPORAL_DELIMITER => { - // Marks the beginning of a new temporal unit. - self.temporal_unit_count += 1; - self.frame_header_parsed_in_tu = false; - } - OBU_SEQUENCE_HEADER => { - self.parse_sequence_header(&obu.payload)?; - } - OBU_FRAME_HEADER | OBU_FRAME => { - if !self.frame_header_parsed_in_tu { - let info = self.parse_frame_display_info(&obu.payload)?; - let tu_idx = self.temporal_unit_count.saturating_sub(1); - self.temporal_units.push(TemporalUnitInfo { - index: tu_idx, - is_displayed: info.0, - is_show_existing: info.1, - }); - self.frame_header_parsed_in_tu = true; - } - } - OBU_REDUNDANT_FRAME_HEADER => { - // Redundant copies carry the same info — skip to avoid duplicates. - } - _ => {} - } - Ok(()) - } - - /// Return all collected temporal unit infos. - pub fn temporal_units(&self) -> &[TemporalUnitInfo] { - &self.temporal_units - } - - /// Returns the number of temporal units that produce a displayed frame. - pub fn display_frame_count(&self) -> usize { - self.temporal_units.iter().filter(|t| t.is_displayed).count() - } - - // ----------------------------------------------------------------------- - // Sequence header parsing - // ----------------------------------------------------------------------- - - fn parse_sequence_header(&mut self, payload: &[u8]) -> Result<()> { - if payload.len() < 1 { - return Ok(()); - } - let mut r = BsIoSliceReader::from_slice(payload); - - // seq_profile (3 bits) - let _seq_profile = r.read::<3, u8>()?; - // still_picture (1 bit) - let _still_picture = r.read::<1, u8>()?; - // reduced_still_picture_header (1 bit) - self.reduced_still_picture_header = r.read::<1, u8>()? != 0; - - Ok(()) - } - - // ----------------------------------------------------------------------- - // Frame header parsing (minimal — only what is needed for display flags) - // ----------------------------------------------------------------------- - - /// Returns `(is_displayed, is_show_existing)`. - fn parse_frame_display_info(&self, payload: &[u8]) -> Result<(bool, bool)> { - if self.reduced_still_picture_header { - // Section 5.9.2: reduced_still_picture_header implies - // show_existing_frame = 0, show_frame = 1. - return Ok((true, false)); - } - - if payload.is_empty() { - return Ok((true, false)); - } - - let mut r = BsIoSliceReader::from_slice(payload); - - // show_existing_frame (f(1)) - let show_existing_frame = r.read::<1, u8>()? != 0; - if show_existing_frame { - return Ok((true, true)); - } - - // frame_type (f(2)) - let _frame_type = r.read::<2, u8>()?; - - // show_frame (f(1)) - let show_frame = r.read::<1, u8>()? != 0; - - if show_frame { - Ok((true, false)) - } else { - // showable_frame (f(1)) - let showable_frame = r.read::<1, u8>()? != 0; - Ok((showable_frame, false)) - } - } -}