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..ffe0324 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::{ + 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,135 @@ 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 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 { + 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)) => { + 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..6a9e6e8 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::{ @@ -12,13 +12,26 @@ 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}; +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 +109,234 @@ 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<()> { + 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." + ); + } - injector.process_input()?; - injector.interleave_hdr10plus_nals() + 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)?; + + 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 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) => { + 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: &[Obu], encoded: &[u8], validate: bool) -> Vec { + 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..176707f 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; @@ -11,8 +11,19 @@ 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 { + matches!( + path.extension().and_then(|e| e.to_str()), + Some("av1") | Some("ivf") + ) +} + pub struct Remover { input: PathBuf, progress_bar: ProgressBar, @@ -20,7 +31,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 +39,99 @@ 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"); - } + 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 hevc_out = match output { + Some(path) => path, + None => PathBuf::from("hdr10plus_removed_output.hevc"), + }; - let pb = initialize_progress_bar(&format, &input)?; + let pb = initialize_progress_bar(&format, &input)?; - let mut remover = Remover { - input, - progress_bar: pb, - writer: BufWriter::with_capacity( - 100_000, - File::create(hevc_out).expect("Can't create file"), - ), - }; + 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) + } + } + + fn remove_sei_av1( + input: PathBuf, + output: Option, + options: CliOptions, + ) -> Result<()> { + let out_path = output.unwrap_or_else(|| PathBuf::from("hdr10plus_removed_output.av1")); + + 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; + + 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), + } + } + } - remover.process_input(&format) + 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..bc0e0d1 --- /dev/null +++ b/src/core/av1_parser.rs @@ -0,0 +1,95 @@ +// Re-export generic types/functions from the crate +pub use av1_parser::{ + IvfFrameHeader, Obu, + OBU_TEMPORAL_DELIMITER, OBU_METADATA, + decode_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() +} 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)]