Couple fixes, RPU parsing inconsistencies

Add progress bar for frame indices computation in `inject-rpu` (closes #27)
Fix `inject-rpu` `rpu-in` argument to match documentation (closes #28)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dovi_tool"
-version = "0.3.4"
+version = "0.3.5"
 authors = ["quietvoid"]
 edition = "2018"
 license = "MIT"
@@ -13,6 +13,6 @@ ansi_term = "0.12.1"
 crc = "2.0.0-rc.1"
 bitvec = "0.20.1"
 bitvec_helpers = "0.1.0"
-hevc_parser = "0.1.4"
+hevc_parser = "0.1.5"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

README.md

@@ -35,7 +35,7 @@ Input can be piped.
 #### inject-rpu
 Interleaves RPU NAL units between slices in an encoded HEVC file.
 
-* `dovi_tool inject-rpu -i video.hevc --rpu-in RPU.bin`
+* `dovi_tool inject-rpu -i video.hevc --rpu-in RPU.bin -o injected_output.hevc`
 
 #### editor
 Edits a RPU according to a JSON config.  

src/commands.rs

@@ -119,15 +119,15 @@ pub enum Command {
         )]
         input: PathBuf,
 
+        #[structopt(long, help = "Sets the input RPU file to use", parse(from_os_str))]
+        rpu_in: PathBuf,
+
         #[structopt(
-            name = "rpu_in",
             long,
-            help = "Sets the input RPU file to use",
+            short = "o",
+            help = "Output HEVC file location",
             parse(from_os_str)
         )]
-        rpu_in: PathBuf,
-
-        #[structopt(long, help = "Output HEVC file location", parse(from_os_str))]
         output: Option<PathBuf>,
     },
 

src/dovi/editor.rs

@@ -101,6 +101,8 @@ impl Editor {
                 config.duplicate_metadata(to_duplicate, &mut data);
             }
 
+            println!("Final metadata length: {}", data.len());
+
             match write_rpu_file(&editor.rpu_out, data) {
                 Ok(_) => (),
                 Err(e) => panic!("{:?}", e),
@@ -191,8 +193,6 @@ impl EditConfig {
                 std::iter::repeat(source).take(meta.length),
             );
         });
-
-        println!("Duplicated, new metadata len {}", data.len());
     }
 }
 

src/dovi/mod.rs

@@ -10,7 +10,7 @@ mod rpu;
 
 use hevc_parser::{
     hevc::{Frame, NAL_AUD},
-    HevcParser,
+    HevcParser, NALUStartCode,
 };
 use rpu::{parse_dovi_rpu, DoviRpu};
 
@@ -89,7 +89,7 @@ pub fn parse_rpu_file(input: &Path) -> Option<Vec<DoviRpu>> {
     reader.read_exact(&mut data).unwrap();
 
     let mut offsets = Vec::with_capacity(200_000);
-    let mut parser = HevcParser::default();
+    let mut parser = HevcParser::with_nalu_start_code(NALUStartCode::Length4);
 
     parser.get_offsets(&data, &mut offsets);
 
@@ -101,21 +101,12 @@ pub fn parse_rpu_file(input: &Path) -> Option<Vec<DoviRpu>> {
         .enumerate()
         .map(|(index, offset)| {
             let size = if offset == &last {
-                data.len() - offset - 1
+                data.len() - offset
             } else {
-                let size = if index == count - 1 {
-                    last - offset
-                } else {
-                    offsets[index + 1] - offset
-                };
-
-                match &data[offset + size - 1..offset + size + 3] {
-                    [0, 0, 0, 1] => size - 2,
-                    _ => size,
-                }
+                offsets[index + 1] - offset
             };
 
-            let start = *offset + 1;
+            let start = *offset;
             let end = start + size;
 
             parse_dovi_rpu(&data[start..end])

src/dovi/rpu/mod.rs

@@ -15,8 +15,22 @@ use hevc_parser::utils::{
 
 #[inline(always)]
 pub fn parse_dovi_rpu(data: &[u8]) -> Result<DoviRpu, String> {
+    if data.len() < 25 {
+        return Err(format!("Invalid RPU\n{:?}", &data));
+    }
+
+    // Including 0x7C01 prepended
+    let trimmed_data = match &data[..5] {
+        [0, 0, 0, 1, 25] => &data[4..],
+        [0, 0, 1, 25, 8] => &data[3..],
+        [0, 1, 25, 8, 9] | [124, 1, 25, 8, 9] => &data[2..],
+        [1, 25, 8, 9, _] => &data[1..],
+        [25, 8, 9, _, _] => &data,
+        _ => return Err(format!("Invalid RPU data start bytes\n{:?}", &data)),
+    };
+
     // Clear start code emulation prevention 3 byte
-    let bytes: Vec<u8> = clear_start_code_emulation_prevention_3_byte(&data[2..]);
+    let bytes: Vec<u8> = clear_start_code_emulation_prevention_3_byte(trimmed_data);
 
     let len = bytes.len();
 

src/dovi/rpu/tests.rs

@@ -180,3 +180,18 @@ fn sets_offsets_to_zero() {
         assert_eq!(vec![0, 0, 0, 0], block._get_offsets());
     }
 }
+
+#[test]
+fn profile8_001_end_crc32() {
+    use crate::dovi::parse_rpu_file;
+
+    let rpus = parse_rpu_file(&PathBuf::from("./assets/p8_001_end_crc32.bin"));
+    assert!(rpus.is_some());
+
+    let rpus = rpus.unwrap();
+    assert_eq!(rpus.len(), 3);
+
+    let dovi_rpu = &rpus[0];
+    assert_eq!(8, dovi_rpu.dovi_profile);
+    assert_eq!([216, 0, 0, 1], dovi_rpu.rpu_data_crc32.to_be_bytes());
+}

src/dovi/rpu_injector.rs

@@ -8,6 +8,7 @@ use super::{input_format, parse_rpu_file, DoviRpu, Format, OUT_NAL_HEADER};
 
 use hevc_parser::hevc::*;
 use hevc_parser::HevcParser;
+use indicatif::{ProgressBar, ProgressStyle};
 
 pub struct RpuInjector {
     input: PathBuf,
@@ -143,11 +144,25 @@ impl RpuInjector {
             println!("Computing frame indices..");
             stdout().flush().ok();
 
+            let pb_indices = ProgressBar::new(frames.len() as u64);
+            pb_indices.set_style(
+                ProgressStyle::default_bar()
+                    .template("[{elapsed_precise}] {bar:60.cyan} {percent}%"),
+            );
+
             let last_slice_indices: Vec<usize> = frames
                 .iter()
-                .map(|f| find_last_slice_nal_index(nals, f))
+                .map(|f| {
+                    let index = find_last_slice_nal_index(nals, f);
+
+                    pb_indices.inc(1);
+
+                    index
+                })
                 .collect();
 
+            pb_indices.finish_and_clear();
+
             assert_eq!(frames.len(), last_slice_indices.len());
 
             println!("Rewriting file with interleaved RPU NALs..");