extract-rpu: Add limit option

quietvoid · Sep 14, 2024 · 1773a4d · 1773a4d
1 parent 46662ed
commit 1773a4d
Show file tree

Hide file tree

Showing 10 changed files with 194 additions and 152 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -14,28 +14,28 @@ path = "src/main.rs"
 
 [dependencies]
 dolby_vision = { path = "dolby_vision", "features" = ["xml", "serde"] }
-bitvec_helpers = { version = "3.1.5", default-features = false, features = ["bitstream-io"] }
-hevc_parser = { version = "0.6.3", features = ["hevc_io"] }
+bitvec_helpers = { version = "3.1.6", default-features = false, features = ["bitstream-io"] }
+hevc_parser = { version = "0.6.4", features = ["hevc_io"] }
 madvr_parse = "1.0.2"
 hdr10plus = { version = "2.1.1", features = ["json"] }
 
-anyhow = "1.0.86"
-clap = { version = "4.5.16", features = ["derive", "wrap_help", "deprecated"] }
+anyhow = "1.0.88"
+clap = { version = "4.5.17", features = ["derive", "wrap_help", "deprecated"] }
 clap_lex = "*"
 indicatif = "0.17.8"
 bitvec = "1.0.1"
-serde = { version = "1.0.208", features = ["derive"] }
-serde_json = { version = "1.0.125", features = ["preserve_order"] }
+serde = { version = "1.0.210", features = ["derive"] }
+serde_json = { version = "1.0.128", features = ["preserve_order"] }
 itertools = "0.13.0"
-plotters = { version = "0.3.6", default-features = false, features = ["bitmap_backend", "bitmap_encoder", "all_series"] }
+plotters = { version = "0.3.7", default-features = false, features = ["bitmap_backend", "bitmap_encoder", "all_series"] }
 
 [dev-dependencies]
 assert_cmd = "2.0.16"
 assert_fs = "1.1.2"
 predicates = "3.1.2"
 
 [build-dependencies]
-anyhow = "1.0.86"
+anyhow = "1.0.88"
 vergen-gitcl = { version = "1.0.0", default-features = false, features = ["build"] }
 
 [features]

diff --git a/README.md b/README.md
@@ -257,6 +257,9 @@ For working with an HEVC source file, there are multiple options that apply to m
  
     **Supports profiles 4, 5, 7, and 8**.
 
+    **Flags**:
+    - `-l`, `--limit` Number of frames to process from the input. Processing stops after N frames.
+
     **Examples**:
     ```console
     dovi_tool extract-rpu video.hevc

diff --git a/src/commands/extract_rpu.rs b/src/commands/extract_rpu.rs
@@ -30,4 +30,12 @@ pub struct ExtractRpuArgs {
         value_hint = ValueHint::FilePath
     )]
     pub rpu_out: Option<PathBuf>,
+
+    #[arg(
+        id = "limit",
+        long,
+        short = 'l',
+        help = "Stop processing input after N frames"
+    )]
+    pub limit: Option<u64>,
 }
diff --git a/src/dovi/converter.rs b/src/dovi/converter.rs
@@ -59,7 +59,13 @@ impl Converter {
 
     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);
+        let mut dovi_processor = DoviProcessor::new(
+            options,
+            self.input.clone(),
+            dovi_writer,
+            pb,
+            Default::default(),
+        );
 
         dovi_processor.read_write_from_io(&self.format)
     }

diff --git a/src/dovi/demuxer.rs b/src/dovi/demuxer.rs
@@ -70,7 +70,13 @@ impl Demuxer {
         };
 
         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);
+        let mut dovi_processor = DoviProcessor::new(
+            options,
+            self.input.clone(),
+            dovi_writer,
+            pb,
+            Default::default(),
+        );
 
         dovi_processor.read_write_from_io(&self.format)
     }

diff --git a/src/dovi/general_read_write.rs b/src/dovi/general_read_write.rs
@@ -24,6 +24,8 @@ pub struct DoviProcessor {
 
     progress_bar: ProgressBar,
     dovi_writer: DoviWriter,
+
+    processor_opts: DoviProcessorOptions,
 }
 
 pub struct DoviWriter {
@@ -40,6 +42,11 @@ pub struct RpuNal {
     data: Vec<u8>,
 }
 
+#[derive(Default)]
+pub struct DoviProcessorOptions {
+    pub limit: Option<u64>,
+}
+
 impl DoviWriter {
     pub fn new<P: AsRef<Path>>(
         bl_out: Option<P>,
@@ -91,6 +98,7 @@ impl DoviProcessor {
         input: PathBuf,
         dovi_writer: DoviWriter,
         progress_bar: ProgressBar,
+        processor_opts: DoviProcessorOptions,
     ) -> DoviProcessor {
         DoviProcessor {
             input,
@@ -101,6 +109,7 @@ impl DoviProcessor {
             previous_rpu_index: 0,
             progress_bar,
             dovi_writer,
+            processor_opts,
         }
     }
 
@@ -109,6 +118,7 @@ impl DoviProcessor {
 
         let processor_opts = HevcProcessorOpts {
             parse_nals: true,
+            limit: self.processor_opts.limit,
             ..Default::default()
         };
         let mut processor = HevcProcessor::new(format.clone(), processor_opts, chunk_size);

diff --git a/src/dovi/muxer.rs b/src/dovi/muxer.rs
@@ -38,7 +38,7 @@ pub struct Muxer {
 pub struct ElHandler {
     input: PathBuf,
     writer: BufWriter<File>,
-    buffers: VecDeque<FrameBuffer>,
+    buffered_frames: VecDeque<FrameBuffer>,
 
     options: CliOptions,
 }
@@ -92,7 +92,7 @@ impl Muxer {
         let el_handler = ElHandler {
             input: el,
             writer,
-            buffers: VecDeque::new(),
+            buffered_frames: VecDeque::new(),
             options: cli_options.clone(),
         };
 
@@ -207,13 +207,13 @@ impl IoProcessor for Muxer {
                 self.write_bl_frame()?;
 
                 // Process EL, read if possibly incomplete frame
-                if self.el_handler.buffers.len() < 2 {
+                if self.el_handler.buffered_frames.len() < 2 {
                     self.el_processor
                         .parse_nalus(&mut self.el_reader, &mut self.el_handler)?;
                 }
 
                 // Write EL frame if complete
-                if self.el_handler.buffers.len() > 1 {
+                if self.el_handler.buffered_frames.len() > 1 {
                     self.el_handler.write_next_frame()?;
                 }
 
@@ -283,14 +283,14 @@ impl IoProcessor for Muxer {
             // Write last BL frame
             self.write_bl_frame()?;
 
-            if self.el_handler.buffers.len() == 1 {
+            if self.el_handler.buffered_frames.len() == 1 {
                 // Maybe incomplete last frame
                 self.el_processor
                     .parse_nalus(&mut self.el_reader, &mut self.el_handler)?;
 
                 // Write last EL frame
                 self.el_handler.write_next_frame()?;
-            } else if let Some(last_frame) = self.el_handler.buffers.back() {
+            } else if let Some(last_frame) = self.el_handler.buffered_frames.back() {
                 // Zero indexed
                 bail!(
                     "Mismatched BL/EL frame count. Expected {} frames, got {} (or more) frames in EL",
@@ -354,7 +354,7 @@ impl IoProcessor for ElHandler {
 
             // Existing incomplete frame
             let existing_frame = self
-                .buffers
+                .buffered_frames
                 .iter_mut()
                 .find(|fb| fb.frame_number == frame_number);
 
@@ -366,7 +366,7 @@ impl IoProcessor for ElHandler {
                     nals: nal_buffers.collect(),
                 };
 
-                self.buffers.push_back(frame_buffer);
+                self.buffered_frames.push_back(frame_buffer);
             }
         }
 
@@ -439,7 +439,7 @@ impl Muxer {
 
 impl ElHandler {
     fn write_next_frame(&mut self) -> Result<()> {
-        if let Some(frame_buffer) = self.buffers.pop_front() {
+        if let Some(frame_buffer) = self.buffered_frames.pop_front() {
             for nal_buf in frame_buffer.nals {
                 let nal_type = if nal_buf.nal_type != NAL_UNSPEC62 {
                     NAL_UNSPEC63

diff --git a/src/dovi/remover.rs b/src/dovi/remover.rs
@@ -52,7 +52,13 @@ impl Remover {
         let bl_out = Some(self.output.as_path());
 
         let dovi_writer = DoviWriter::new(bl_out, None, None, None);
-        let mut dovi_processor = DoviProcessor::new(options, self.input.clone(), dovi_writer, pb);
+        let mut dovi_processor = DoviProcessor::new(
+            options,
+            self.input.clone(),
+            dovi_writer,
+            pb,
+            Default::default(),
+        );
 
         dovi_processor.read_write_from_io(&self.format)
     }

diff --git a/src/dovi/rpu_extractor.rs b/src/dovi/rpu_extractor.rs
@@ -4,13 +4,17 @@ use std::path::PathBuf;
 
 use crate::commands::ExtractRpuArgs;
 
-use super::{general_read_write, input_from_either, CliOptions, IoFormat};
+use super::{
+    general_read_write::{self, DoviProcessorOptions},
+    input_from_either, CliOptions, IoFormat,
+};
 use general_read_write::{DoviProcessor, DoviWriter};
 
 pub struct RpuExtractor {
     format: IoFormat,
     input: PathBuf,
     rpu_out: PathBuf,
+    limit: Option<u64>,
 }
 
 impl RpuExtractor {
@@ -19,6 +23,7 @@ impl RpuExtractor {
             input,
             input_pos,
             rpu_out,
+            limit,
         } = args;
 
         let input = input_from_either("extract-rpu", input, input_pos)?;
@@ -33,6 +38,7 @@ impl RpuExtractor {
             format,
             input,
             rpu_out,
+            limit,
         })
     }
 
@@ -52,7 +58,13 @@ impl RpuExtractor {
 
     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);
+        let mut dovi_processor = DoviProcessor::new(
+            options,
+            self.input.clone(),
+            dovi_writer,
+            pb,
+            DoviProcessorOptions { limit: self.limit },
+        );
 
         dovi_processor.read_write_from_io(&self.format)
     }