From d6bed49bf8d066695e3864121eeaa662a986c6f0 Mon Sep 17 00:00:00 2001 From: quietvoid <39477805+quietvoid@users.noreply.github.com> Date: Sun, 1 May 2022 01:33:34 -0400 Subject: [PATCH] Refactor generic read/write to use new I/O --- Cargo.lock | 4 +- Cargo.toml | 2 +- src/dovi/converter.rs | 12 +-- src/dovi/demuxer.rs | 13 ++- src/dovi/general_read_write.rs | 190 ++++++++++++--------------------- src/dovi/muxer.rs | 9 +- src/dovi/rpu_extractor.rs | 12 +-- src/dovi/rpu_injector.rs | 3 +- 8 files changed, 95 insertions(+), 150 deletions(-) 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 b04727c..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) } } 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 a4a8934..0d78db4 100644 --- a/src/dovi/rpu_injector.rs +++ b/src/dovi/rpu_injector.rs @@ -335,7 +335,8 @@ fn find_last_slice_nal_index(nals: &[NALUnit], frame: &Frame) -> usize { let last_slice_nal = last_slice.1 .1; // Use last non EOS/EOB NALU - let non_eos_eob_nal_count = frame.nals + let non_eos_eob_nal_count = frame + .nals .iter() .filter(|nal| !matches!(nal.nal_type, NAL_EOS_NUT | NAL_EOB_NUT)) .count();