diff --git a/Cargo.lock b/Cargo.lock index aa42f2f..e30db1e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -262,9 +262,9 @@ dependencies = [ [[package]] name = "hevc_parser" -version = "0.4.2" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "734c1edab5d988dd78c7848bdf478388d57a3e129d2710d5e402c8712891b38d" +checksum = "f3ccf98b57f0328d11af098ef87ee6742edc59b2abdbc2fd38189323f4396b0d" dependencies = [ "anyhow", "bitvec_helpers", diff --git a/Cargo.toml b/Cargo.toml index d398f09..25c34fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,7 +8,7 @@ license = "MIT" [dependencies] bitvec_helpers = "1.0.2" -hevc_parser = { version = "0.4.2", features = ["hevc_io"] } +hevc_parser = { version = "0.4.3", features = ["hevc_io"] } dolby_vision = { path = "dolby_vision", "features" = ["xml", "serde_feature"] } madvr_parse = { path = "madvr_parse" } diff --git a/src/dovi/converter.rs b/src/dovi/converter.rs index a04c60e..ddf87c6 100644 --- a/src/dovi/converter.rs +++ b/src/dovi/converter.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use super::{general_read_write, CliOptions, IoFormat}; -use general_read_write::{DoviReader, DoviWriter}; +use general_read_write::{DoviProcessor, DoviWriter}; pub struct Converter { format: IoFormat, @@ -54,14 +54,14 @@ impl Converter { match self.format { IoFormat::Matroska => bail!("Converter: Matroska input is unsupported"), - _ => self.convert_raw_hevc(Some(&pb), options), + _ => self.convert_raw_hevc(pb, options), } } - fn convert_raw_hevc(&self, pb: Option<&ProgressBar>, options: CliOptions) -> Result<()> { - let mut dovi_reader = DoviReader::new(options); - let mut dovi_writer = DoviWriter::new(None, None, None, Some(&self.output)); + 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(options, self.input.clone(), dovi_writer, pb); - dovi_reader.read_write_from_io(&self.format, &self.input, pb, &mut dovi_writer) + dovi_processor.read_write_from_io(&self.format) } } diff --git a/src/dovi/demuxer.rs b/src/dovi/demuxer.rs index e92f275..036a30a 100644 --- a/src/dovi/demuxer.rs +++ b/src/dovi/demuxer.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use super::{general_read_write, CliOptions, IoFormat}; -use general_read_write::{DoviReader, DoviWriter}; +use general_read_write::{DoviProcessor, DoviWriter}; pub struct Demuxer { format: IoFormat, @@ -68,21 +68,20 @@ impl Demuxer { match self.format { IoFormat::Matroska => bail!("Demuxer: Matroska input is unsupported"), - _ => self.demux_raw_hevc(Some(&pb), options), + _ => self.demux_raw_hevc(pb, options), } } - fn demux_raw_hevc(&self, pb: Option<&ProgressBar>, options: CliOptions) -> Result<()> { - let mut dovi_reader = DoviReader::new(options); - + fn demux_raw_hevc(&self, pb: ProgressBar, options: CliOptions) -> Result<()> { let bl_out = if self.el_only { None } else { Some(self.bl_out.as_path()) }; - let mut dovi_writer = DoviWriter::new(bl_out, Some(self.el_out.as_path()), None, None); + let dovi_writer = DoviWriter::new(bl_out, Some(self.el_out.as_path()), None, None); + let mut dovi_processor = DoviProcessor::new(options, self.input.clone(), dovi_writer, pb); - dovi_reader.read_write_from_io(&self.format, &self.input, pb, &mut dovi_writer) + dovi_processor.read_write_from_io(&self.format) } } diff --git a/src/dovi/general_read_write.rs b/src/dovi/general_read_write.rs index 75b1432..de410d9 100644 --- a/src/dovi/general_read_write.rs +++ b/src/dovi/general_read_write.rs @@ -1,20 +1,27 @@ -use anyhow::{bail, Result}; -use indicatif::ProgressBar; -use std::io::Read; use std::io::{stdout, BufRead, BufReader, BufWriter, Write}; +use std::path::PathBuf; use std::{fs::File, path::Path}; +use anyhow::{bail, Result}; +use indicatif::ProgressBar; + use hevc_parser::hevc::NALUnit; use hevc_parser::hevc::{NAL_SEI_PREFIX, NAL_UNSPEC62, NAL_UNSPEC63}; +use hevc_parser::io::{processor, IoProcessor}; use hevc_parser::HevcParser; +use processor::{HevcProcessor, HevcProcessorOpts}; use super::{convert_encoded_from_opts, is_st2094_40_sei, CliOptions, IoFormat, OUT_NAL_HEADER}; -pub struct DoviReader { +pub struct DoviProcessor { + input: PathBuf, options: CliOptions, rpu_nals: Vec, previous_rpu_index: u64, + + progress_bar: ProgressBar, + dovi_writer: DoviWriter, } pub struct DoviWriter { @@ -76,128 +83,46 @@ impl DoviWriter { } } -impl DoviReader { - pub fn new(options: CliOptions) -> DoviReader { - DoviReader { +impl DoviProcessor { + pub fn new( + options: CliOptions, + input: PathBuf, + dovi_writer: DoviWriter, + progress_bar: ProgressBar, + ) -> DoviProcessor { + DoviProcessor { + input, options, rpu_nals: Vec::new(), previous_rpu_index: 0, + progress_bar, + dovi_writer, } } - pub fn read_write_from_io( - &mut self, - format: &IoFormat, - input: &Path, - pb: Option<&ProgressBar>, - dovi_writer: &mut DoviWriter, - ) -> Result<()> { - //BufReader & BufWriter - let stdin = std::io::stdin(); - let mut reader = Box::new(stdin.lock()) as Box; - - if let IoFormat::Raw = format { - let file = File::open(input)?; - reader = Box::new(BufReader::with_capacity(100_000, file)); - } - + pub fn read_write_from_io(&mut self, format: &IoFormat) -> Result<()> { let chunk_size = 100_000; - let mut main_buf = vec![0; 100_000]; - let mut sec_buf = vec![0; 50_000]; - - let mut chunk = Vec::with_capacity(chunk_size); - let mut end: Vec = Vec::with_capacity(100_000); - - let mut consumed = 0; - - let mut parser = HevcParser::default(); - - let mut offsets = Vec::with_capacity(2048); - let parse_nals = dovi_writer.rpu_writer.is_some(); - - while let Ok(n) = reader.read(&mut main_buf) { - let mut read_bytes = n; - if read_bytes == 0 && end.is_empty() && chunk.is_empty() { - break; - } - - if *format == IoFormat::RawStdin { - chunk.extend_from_slice(&main_buf[..read_bytes]); - - loop { - let num = reader.read(&mut sec_buf)?; - - if num > 0 { - read_bytes += num; - - chunk.extend_from_slice(&sec_buf[..num]); - - if read_bytes >= chunk_size { - break; - } - } else { - break; - } - } - } else if read_bytes < chunk_size { - chunk.extend_from_slice(&main_buf[..read_bytes]); - } else { - chunk.extend_from_slice(&main_buf); - } + let parse_nals = self.dovi_writer.rpu_writer.is_some(); - parser.get_offsets(&chunk, &mut offsets); + let processor_opts = HevcProcessorOpts { + parse_nals, + ..Default::default() + }; + let mut processor = HevcProcessor::new(format.clone(), processor_opts, chunk_size); - if offsets.is_empty() { - continue; - } - - let last = if read_bytes < chunk_size { - *offsets.last().unwrap() - } else { - let last = offsets.pop().unwrap(); - - end.clear(); - end.extend_from_slice(&chunk[last..]); - - last - }; - - let nals: Vec = parser.split_nals(&chunk, &offsets, last, parse_nals)?; - self.write_nals(&chunk, dovi_writer, &nals)?; - - chunk.clear(); - - if !end.is_empty() { - chunk.extend_from_slice(&end); - end.clear(); - } - - consumed += read_bytes; - - if consumed >= 100_000_000 { - if let Some(pb) = pb { - pb.inc(1); - consumed = 0; - } - } - } + let stdin = std::io::stdin(); + let mut reader = Box::new(stdin.lock()) as Box; - if let Some(pb) = pb { - pb.finish_and_clear(); + if let IoFormat::Raw = format { + let file = File::open(&self.input)?; + reader = Box::new(BufReader::with_capacity(100_000, file)); } - parser.finish(); - - self.flush_writer(&parser, dovi_writer) + processor.process_io(&mut reader, self) } - pub fn write_nals( - &mut self, - chunk: &[u8], - dovi_writer: &mut DoviWriter, - nals: &[NALUnit], - ) -> Result<()> { + pub fn write_nals(&mut self, chunk: &[u8], nals: &[NALUnit]) -> Result<()> { for nal in nals { if self.options.drop_hdr10plus && nal.nal_type == NAL_SEI_PREFIX @@ -220,7 +145,7 @@ impl DoviReader { continue; } - if let Some(ref mut sl_writer) = dovi_writer.sl_writer { + if let Some(ref mut sl_writer) = self.dovi_writer.sl_writer { if nal.nal_type == NAL_UNSPEC63 && self.options.discard_el { continue; } @@ -245,7 +170,7 @@ impl DoviReader { match nal.nal_type { NAL_UNSPEC63 => { - if let Some(ref mut el_writer) = dovi_writer.el_writer { + if let Some(ref mut el_writer) = self.dovi_writer.el_writer { el_writer.write_all(OUT_NAL_HEADER)?; el_writer.write_all(&chunk[nal.start + 2..nal.end])?; } @@ -253,7 +178,7 @@ impl DoviReader { NAL_UNSPEC62 => { self.previous_rpu_index = nal.decoded_frame_index; - if let Some(ref mut el_writer) = dovi_writer.el_writer { + if let Some(ref mut el_writer) = self.dovi_writer.el_writer { el_writer.write_all(OUT_NAL_HEADER)?; } @@ -267,29 +192,29 @@ impl DoviReader { if let Some(_mode) = self.options.mode { let modified_data = convert_encoded_from_opts(&self.options, rpu_data)?; - if let Some(ref mut _rpu_writer) = dovi_writer.rpu_writer { + if let Some(ref mut _rpu_writer) = self.dovi_writer.rpu_writer { // RPU for x265, remove 0x7C01 self.rpu_nals.push(RpuNal { decoded_index: self.rpu_nals.len(), presentation_number: 0, data: modified_data[2..].to_owned(), }); - } else if let Some(ref mut el_writer) = dovi_writer.el_writer { + } else if let Some(ref mut el_writer) = self.dovi_writer.el_writer { el_writer.write_all(&modified_data)?; } - } else if let Some(ref mut _rpu_writer) = dovi_writer.rpu_writer { + } else if let Some(ref mut _rpu_writer) = self.dovi_writer.rpu_writer { // RPU for x265, remove 0x7C01 self.rpu_nals.push(RpuNal { decoded_index: self.rpu_nals.len(), presentation_number: 0, data: rpu_data[2..].to_vec(), }); - } else if let Some(ref mut el_writer) = dovi_writer.el_writer { + } else if let Some(ref mut el_writer) = self.dovi_writer.el_writer { el_writer.write_all(rpu_data)?; } } _ => { - if let Some(ref mut bl_writer) = dovi_writer.bl_writer { + if let Some(ref mut bl_writer) = self.dovi_writer.bl_writer { bl_writer.write_all(OUT_NAL_HEADER)?; bl_writer.write_all(&chunk[nal.start..nal.end])?; } @@ -300,17 +225,17 @@ impl DoviReader { Ok(()) } - fn flush_writer(&mut self, parser: &HevcParser, dovi_writer: &mut DoviWriter) -> Result<()> { - if let Some(ref mut bl_writer) = dovi_writer.bl_writer { + fn flush_writer(&mut self, parser: &HevcParser) -> Result<()> { + if let Some(ref mut bl_writer) = self.dovi_writer.bl_writer { bl_writer.flush()?; } - if let Some(ref mut el_writer) = dovi_writer.el_writer { + if let Some(ref mut el_writer) = self.dovi_writer.el_writer { el_writer.flush()?; } // Reorder RPUs to display output order - if let Some(ref mut rpu_writer) = dovi_writer.rpu_writer { + if let Some(ref mut rpu_writer) = self.dovi_writer.rpu_writer { let frames = parser.ordered_frames(); if frames.is_empty() { @@ -356,3 +281,22 @@ impl DoviReader { Ok(()) } } + +impl IoProcessor for DoviProcessor { + fn input(&self) -> &std::path::PathBuf { + &self.input + } + + fn update_progress(&mut self, delta: u64) { + self.progress_bar.inc(delta); + } + + fn process_nals(&mut self, _parser: &HevcParser, nals: &[NALUnit], chunk: &[u8]) -> Result<()> { + self.write_nals(chunk, nals) + } + + fn finalize(&mut self, parser: &HevcParser) -> Result<()> { + self.progress_bar.finish_and_clear(); + self.flush_writer(parser) + } +} diff --git a/src/dovi/muxer.rs b/src/dovi/muxer.rs index 40b8967..92e4b9d 100644 --- a/src/dovi/muxer.rs +++ b/src/dovi/muxer.rs @@ -82,7 +82,10 @@ impl Muxer { let el_file = File::open(&el)?; let el_reader = Box::new(BufReader::with_capacity(chunk_size, el_file)); - let el_opts = HevcProcessorOpts { buffer_frame: true }; + let el_opts = HevcProcessorOpts { + buffer_frame: true, + ..Default::default() + }; let el_handler = ElHandler { input: el, writer, @@ -133,9 +136,7 @@ impl Muxer { reader = Box::new(BufReader::with_capacity(100_000, file)); } - processor.process_io(&mut reader, self)?; - - Ok(()) + processor.process_io(&mut reader, self) } } @@ -270,6 +271,8 @@ impl IoProcessor for Muxer { self.el_handler.writer.flush()?; + self.progress_bar.finish_and_clear(); + Ok(()) } } diff --git a/src/dovi/rpu_extractor.rs b/src/dovi/rpu_extractor.rs index f032003..a1b53e0 100644 --- a/src/dovi/rpu_extractor.rs +++ b/src/dovi/rpu_extractor.rs @@ -3,7 +3,7 @@ use indicatif::ProgressBar; use std::path::PathBuf; use super::{general_read_write, CliOptions, IoFormat}; -use general_read_write::{DoviReader, DoviWriter}; +use general_read_write::{DoviProcessor, DoviWriter}; pub struct RpuExtractor { format: IoFormat, @@ -50,14 +50,14 @@ impl RpuExtractor { match self.format { IoFormat::Matroska => bail!("Extractor: Matroska input is unsupported"), - _ => self.extract_rpu_from_el(Some(&pb), options), + _ => self.extract_rpu_from_el(pb, options), } } - fn extract_rpu_from_el(&self, pb: Option<&ProgressBar>, options: CliOptions) -> Result<()> { - let mut dovi_reader = DoviReader::new(options); - let mut dovi_writer = DoviWriter::new(None, None, Some(&self.rpu_out), None); + fn extract_rpu_from_el(&self, pb: ProgressBar, options: CliOptions) -> Result<()> { + let dovi_writer = DoviWriter::new(None, None, Some(&self.rpu_out), None); + let mut dovi_processor = DoviProcessor::new(options, self.input.clone(), dovi_writer, pb); - dovi_reader.read_write_from_io(&self.format, &self.input, pb, &mut dovi_writer) + dovi_processor.read_write_from_io(&self.format) } } diff --git a/src/dovi/rpu_injector.rs b/src/dovi/rpu_injector.rs index 8f2e2eb..0d78db4 100644 --- a/src/dovi/rpu_injector.rs +++ b/src/dovi/rpu_injector.rs @@ -1,5 +1,5 @@ use std::fs::File; -use std::io::{stdout, BufReader, BufWriter, Read, Write}; +use std::io::{stdout, BufReader, BufWriter, Write}; use std::path::PathBuf; use anyhow::{bail, ensure, Result}; @@ -7,7 +7,9 @@ use indicatif::{ProgressBar, ProgressStyle}; use rayon::prelude::*; use hevc_parser::hevc::*; +use hevc_parser::io::{processor, IoProcessor}; use hevc_parser::HevcParser; +use processor::{HevcProcessor, HevcProcessorOpts}; use super::{ get_aud, is_st2094_40_sei, parse_rpu_file, CliOptions, DoviRpu, IoFormat, OUT_NAL_HEADER, @@ -16,14 +18,67 @@ use super::{ pub struct RpuInjector { input: PathBuf, rpu_in: PathBuf, - output: PathBuf, no_add_aud: bool, options: CliOptions, rpus: Option>, + + writer: BufWriter, + progress_bar: ProgressBar, + already_checked_for_rpu: bool, + + frames: Vec, + nals: Vec, + mismatched_length: bool, + + last_slice_indices: Vec, + last_frame_index: u64, + nals_parsed: usize, + last_metadata_written: Option>, } impl RpuInjector { + pub fn new( + input: PathBuf, + rpu_in: PathBuf, + output: PathBuf, + no_add_aud: bool, + cli_options: CliOptions, + ) -> Result { + let chunk_size = 100_000; + let progress_bar = super::initialize_progress_bar(&IoFormat::Raw, &input)?; + + let writer = BufWriter::with_capacity( + chunk_size, + File::create(&output).expect("Can't create file"), + ); + + let mut injector = RpuInjector { + input, + rpu_in, + no_add_aud, + options: cli_options, + rpus: None, + + writer, + progress_bar, + already_checked_for_rpu: false, + + frames: Vec::new(), + nals: Vec::new(), + mismatched_length: false, + + last_slice_indices: Vec::new(), + last_frame_index: 0, + nals_parsed: 0, + last_metadata_written: None, + }; + + injector.rpus = parse_rpu_file(&injector.rpu_in)?; + + Ok(injector) + } + pub fn inject_rpu( input: PathBuf, rpu_in: PathBuf, @@ -40,139 +95,41 @@ impl RpuInjector { }; let mut injector = RpuInjector::new(input, rpu_in, output, no_add_aud, cli_options)?; - let mut parser = HevcParser::default(); - - injector.process_input(&mut parser, format)?; - parser.finish(); - let frames = parser.ordered_frames(); - let nals = parser.get_nals(); - - injector.interleave_rpu_nals(nals, frames) + injector.process_input()?; + injector.interleave_rpu_nals() } else { bail!("unsupported format") } } - fn process_input(&self, parser: &mut HevcParser, format: IoFormat) -> Result<()> { + fn process_input(&mut self) -> Result<()> { println!("Processing input video for frame order info..."); stdout().flush().ok(); - let pb = super::initialize_progress_bar(&format, &self.input)?; - - //BufReader & BufWriter - let file = File::open(&self.input)?; - let mut reader = Box::new(BufReader::with_capacity(100_000, file)); - let chunk_size = 100_000; - let mut main_buf = vec![0; 100_000]; - - let mut chunk = Vec::with_capacity(chunk_size); - let mut end: Vec = Vec::with_capacity(chunk_size); - - let mut consumed = 0; - - let mut offsets = Vec::with_capacity(2048); - - let mut already_checked_for_rpu = false; - - while let Ok(n) = reader.read(&mut main_buf) { - let read_bytes = n; - if read_bytes == 0 && end.is_empty() && chunk.is_empty() { - break; - } - - if read_bytes < chunk_size { - chunk.extend_from_slice(&main_buf[..read_bytes]); - } else { - chunk.extend_from_slice(&main_buf); - } - - parser.get_offsets(&chunk, &mut offsets); - - if offsets.is_empty() { - continue; - } - - let last = if read_bytes < chunk_size { - *offsets.last().unwrap() - } else { - let last = offsets.pop().unwrap(); - - end.clear(); - end.extend_from_slice(&chunk[last..]); - - last - }; - - if !already_checked_for_rpu { - let nals = parser.split_nals(&chunk, &offsets, last, true)?; - - if nals.iter().any(|e| e.nal_type == NAL_UNSPEC62) { - already_checked_for_rpu = true; - println!("\nWarning: Input file already has RPUs, they will be replaced."); - } - } else { - parser.split_nals(&chunk, &offsets, last, true)?; - } - - chunk.clear(); - - if !end.is_empty() { - chunk.extend_from_slice(&end); - end.clear(); - } - - consumed += read_bytes; - - if consumed >= 100_000_000 { - // Stop checking after 100M - if !already_checked_for_rpu { - already_checked_for_rpu = true; - } + let mut processor = + HevcProcessor::new(IoFormat::Raw, HevcProcessorOpts::default(), chunk_size); - pb.inc(1); - consumed = 0; - } - } + let file = File::open(&self.input)?; + let mut reader = Box::new(BufReader::with_capacity(100_000, file)); - pb.finish_and_clear(); + processor.process_io(&mut reader, self)?; Ok(()) } - pub fn new( - input: PathBuf, - rpu_in: PathBuf, - output: PathBuf, - no_add_aud: bool, - cli_options: CliOptions, - ) -> Result { - let mut injector = RpuInjector { - input, - rpu_in, - output, - no_add_aud, - options: cli_options, - rpus: None, - }; - - injector.rpus = parse_rpu_file(&injector.rpu_in)?; - - Ok(injector) - } - - fn interleave_rpu_nals(&mut self, nals: &[NALUnit], frames: &[Frame]) -> Result<()> { + fn interleave_rpu_nals(&mut self) -> Result<()> { if let Some(ref mut rpus) = self.rpus { - let mismatched_length = if frames.len() != rpus.len() { + self.mismatched_length = if self.frames.len() != rpus.len() { println!( "\nWarning: mismatched lengths. video {}, RPU {}", - frames.len(), + self.frames.len(), rpus.len() ); - if rpus.len() < frames.len() { + if rpus.len() < self.frames.len() { println!("Metadata will be duplicated at the end to match video length\n"); } else { println!("Metadata will be skipped at the end to match video length\n"); @@ -186,16 +143,17 @@ impl RpuInjector { println!("Computing frame indices.."); stdout().flush().ok(); - let pb_indices = ProgressBar::new(frames.len() as u64); + let pb_indices = ProgressBar::new(self.frames.len() as u64); pb_indices.set_style( ProgressStyle::default_bar() .template("[{elapsed_precise}] {bar:60.cyan} {percent}%"), ); - let last_slice_indices: Vec = frames + self.last_slice_indices = self + .frames .par_iter() .map(|f| { - let index = find_last_slice_nal_index(nals, f); + let index = find_last_slice_nal_index(&self.nals, f); pb_indices.inc(1); @@ -205,77 +163,55 @@ impl RpuInjector { pb_indices.finish_and_clear(); - ensure!(frames.len() == last_slice_indices.len()); + ensure!(self.frames.len() == self.last_slice_indices.len()); println!("Rewriting file with interleaved RPU NALs.."); stdout().flush().ok(); - let pb = super::initialize_progress_bar(&IoFormat::Raw, &self.input)?; - let mut parser = HevcParser::default(); + self.progress_bar = super::initialize_progress_bar(&IoFormat::Raw, &self.input)?; let chunk_size = 100_000; - let mut main_buf = vec![0; 100_000]; + let mut processor = + HevcProcessor::new(IoFormat::Raw, HevcProcessorOpts::default(), chunk_size); - let mut chunk = Vec::with_capacity(chunk_size); - let mut end: Vec = Vec::with_capacity(chunk_size); - - //BufReader & BufWriter let file = File::open(&self.input)?; - let mut reader = Box::new(BufReader::with_capacity(100_000, file)); - let mut writer = BufWriter::with_capacity( - chunk_size, - File::create(&self.output).expect("Can't create file"), - ); - - let mut consumed = 0; - let mut offsets = Vec::with_capacity(2048); - - let mut nals_parsed = 0; - - let mut last_metadata_written: Option> = None; - let mut last_frame_index = 0; + let mut reader = Box::new(BufReader::with_capacity(chunk_size, file)); // First frame AUD if !self.no_add_aud { - let first_decoded_frame = frames + let first_decoded_frame = self + .frames .iter() - .find(|f| f.decoded_number == last_frame_index) + .find(|f| f.decoded_number == self.last_frame_index) .unwrap(); - writer.write_all(&get_aud(first_decoded_frame))?; + self.writer.write_all(&get_aud(first_decoded_frame))?; } - while let Ok(n) = reader.read(&mut main_buf) { - let read_bytes = n; - if read_bytes == 0 && end.is_empty() && chunk.is_empty() { - break; - } - - if read_bytes < chunk_size { - chunk.extend_from_slice(&main_buf[..read_bytes]); - } else { - chunk.extend_from_slice(&main_buf); - } - - parser.get_offsets(&chunk, &mut offsets); - - if offsets.is_empty() { - continue; - } + processor.process_io(&mut reader, self)?; + } - let last = if read_bytes < chunk_size { - *offsets.last().unwrap() - } else { - let last = offsets.pop().unwrap(); + Ok(()) + } +} - end.clear(); - end.extend_from_slice(&chunk[last..]); +impl IoProcessor for RpuInjector { + fn input(&self) -> &PathBuf { + &self.input + } - last - }; + fn update_progress(&mut self, delta: u64) { + if !self.already_checked_for_rpu { + self.already_checked_for_rpu = true; + } - let nals = parser.split_nals(&chunk, &offsets, last, true)?; + self.progress_bar.inc(delta); + } + fn process_nals(&mut self, _parser: &HevcParser, nals: &[NALUnit], chunk: &[u8]) -> Result<()> { + // Second pass + if !self.frames.is_empty() && !self.nals.is_empty() { + if let Some(ref mut rpus) = self.rpus { for (cur_index, nal) in nals.iter().enumerate() { // On new frame, write AUD if !self.no_add_aud { @@ -284,14 +220,15 @@ impl RpuInjector { continue; } - if last_frame_index != nal.decoded_frame_index { - let decoded_frame = frames + if self.last_frame_index != nal.decoded_frame_index { + let decoded_frame = self + .frames .iter() .find(|f| f.decoded_number == nal.decoded_frame_index) .unwrap(); - writer.write_all(&get_aud(decoded_frame))?; + self.writer.write_all(&get_aud(decoded_frame))?; - last_frame_index = decoded_frame.decoded_number; + self.last_frame_index = decoded_frame.decoded_number; } } @@ -304,16 +241,17 @@ impl RpuInjector { if nal.nal_type != NAL_UNSPEC62 { // Skip writing existing RPUs, only one allowed - writer.write_all(OUT_NAL_HEADER)?; - writer.write_all(&chunk[nal.start..nal.end])?; + self.writer.write_all(OUT_NAL_HEADER)?; + self.writer.write_all(&chunk[nal.start..nal.end])?; } - let global_index = nals_parsed + cur_index; + let global_index = self.nals_parsed + cur_index; // Slice before interleaved RPU - if last_slice_indices.contains(&global_index) { + if self.last_slice_indices.contains(&global_index) { // We can unwrap because parsed indices are the same - let rpu_index = last_slice_indices + let rpu_index = self + .last_slice_indices .iter() .position(|i| i == &global_index) .unwrap(); @@ -324,43 +262,41 @@ impl RpuInjector { let dovi_rpu = &mut rpus[rpu_index]; let data = dovi_rpu.write_hevc_unspec62_nalu()?; - writer.write_all(OUT_NAL_HEADER)?; - writer.write_all(&data)?; + self.writer.write_all(OUT_NAL_HEADER)?; + self.writer.write_all(&data)?; - last_metadata_written = Some(data); - } else if mismatched_length { - if let Some(data) = &last_metadata_written { - writer.write_all(OUT_NAL_HEADER)?; - writer.write_all(data)?; + self.last_metadata_written = Some(data); + } else if self.mismatched_length { + if let Some(data) = &self.last_metadata_written { + self.writer.write_all(OUT_NAL_HEADER)?; + self.writer.write_all(data)?; } } } } - nals_parsed += nals.len(); - - chunk.clear(); - - if !end.is_empty() { - chunk.extend_from_slice(&end); - end.clear() - } - - consumed += read_bytes; - - if consumed >= 100_000_000 { - pb.inc(1); - consumed = 0; - } + self.nals_parsed += nals.len(); } + } else if !self.already_checked_for_rpu && nals.iter().any(|e| e.nal_type == NAL_UNSPEC62) { + self.already_checked_for_rpu = true; + println!("\nWarning: Input file already has RPUs, they will be replaced."); + } - parser.finish(); - - writer.flush()?; + Ok(()) + } - pb.finish_and_clear(); + fn finalize(&mut self, parser: &HevcParser) -> Result<()> { + // First pass + if self.frames.is_empty() && self.nals.is_empty() { + self.frames = parser.ordered_frames().clone(); + self.nals = parser.get_nals().clone(); + } else { + // Second pass + self.writer.flush()?; } + self.progress_bar.finish_and_clear(); + Ok(()) } } @@ -398,8 +334,14 @@ fn find_last_slice_nal_index(nals: &[NALUnit], frame: &Frame) -> usize { let last_slice_global_index = last_slice.1 .0; let last_slice_nal = last_slice.1 .1; - // Use the last nal because there might be suffix NALs (EL or SEI suffix) - let last_nal_offset = last_slice_index + frame.nals.len() - last_slice_global_index - 1; + // Use last non EOS/EOB NALU + let non_eos_eob_nal_count = frame + .nals + .iter() + .filter(|nal| !matches!(nal.nal_type, NAL_EOS_NUT | NAL_EOB_NUT)) + .count(); + + let last_nal_offset = last_slice_index + non_eos_eob_nal_count - last_slice_global_index - 1; if let Some(first_slice_index) = nals.iter().position(|n| { n.decoded_frame_index == frame.decoded_number && last_slice_nal.nal_type == n.nal_type