Improve export subcommand for different data outputs (#256)

Implements #223

Cargo.toml

@@ -3,7 +3,7 @@ name = "dovi_tool"
 version = "2.0.3"
 authors = ["quietvoid"]
 edition = "2021"
-rust-version = "1.65.0"
+rust-version = "1.70.0"
 license = "MIT"
 build = "build.rs"
 
@@ -18,22 +18,23 @@ hevc_parser = { version = "0.6.1", features = ["hevc_io"] }
 madvr_parse = "1.0.1"
 hdr10plus = { version = "2.1.0", features = ["json"] }
 
-anyhow = "1.0.72"
-clap = { version = "4.3.19", features = ["derive", "wrap_help", "deprecated"] }
-indicatif = "0.17.5"
+anyhow = "1.0.75"
+clap = { version = "4.4.6", features = ["derive", "wrap_help", "deprecated"] }
+clap_lex = "*"
+indicatif = "0.17.7"
 bitvec = "1.0.1"
-serde = { version = "1.0.175", features = ["derive"] }
-serde_json = { version = "1.0.103", features = ["preserve_order"] }
+serde = { version = "1.0.188", features = ["derive"] }
+serde_json = { version = "1.0.107", features = ["preserve_order"] }
 itertools = "0.11.0"
 plotters = { version = "0.3.5", default-features = false, features = ["bitmap_backend", "bitmap_encoder", "all_series"] }
 
 [dev-dependencies]
-assert_cmd = "2.0.11"
+assert_cmd = "2.0.12"
 assert_fs = "1.0.13"
-predicates = "3.0.3"
+predicates = "3.0.4"
 
 [build-dependencies]
-vergen = { version = "8.2.4", default-features = false, features = ["build", "git", "gitcl"] }
+vergen = { version = "8.2.5", default-features = false, features = ["build", "git", "gitcl"] }
 
 [features]
 default = ["system-font"]

README.md

@@ -124,11 +124,25 @@ dovi_tool <SUBCOMMAND> --help
 
 &nbsp;
 * ### **export**
-    Allows exporting a binary RPU file to JSON for simpler analysis.
+    Allows exporting a binary RPU file to text files containing relevant information.  
+    The command allows specifying the desired data to export as file.  
+    **Default**: `export` outputs the full RPU serialized to JSON (equivalent to `--data all`).
 
-    **Example**:
+    * `-d`, `--data`: List of key-value export parameters formatted as `key=output,key2...`
+      * `all` - Exports the list of RPUs as a JSON file
+      * `scenes` - Exports the frame indices at which `scene_refresh_flag` is set to 1
+      * `level5` - Exports the video's L5 metadata in the form of an `editor` config JSON
+
+    &nbsp;
+
+    **Example to export the whole RPU list to JSON**:
+    ```console
+    dovi_tool export -i RPU.bin -d all=RPU_export.json
+    ```
+
+    **Example to export both scene change frames and L5 metadata (with specific path)**
     ```console
-    dovi_tool export -i RPU.bin -o RPU_export.json
+    dovi_tool export -i RPU.bin -d scenes,level5=L5.json
     ```
 
 &nbsp;

dolby_vision/Cargo.toml

@@ -10,12 +10,12 @@ repository = "https://github.com/quietvoid/dovi_tool/tree/main/dolby_vision"
 
 [dependencies]
 bitvec_helpers = { version = "3.1.2", default-features = false, features = ["bitstream-io"] }
-anyhow = "1.0.72"
+anyhow = "1.0.75"
 bitvec = "1.0.1"
 crc = "3.0.1"
-serde = { version = "1.0.175", features = ["derive"], "optional" = true }
-serde_json = { version = "1.0.103", features = ["preserve_order"], "optional" = true }
-roxmltree = { version = "0.18.0", optional = true }
+serde = { version = "1.0.188", features = ["derive"], "optional" = true }
+serde_json = { version = "1.0.107", features = ["preserve_order"], "optional" = true }
+roxmltree = { version = "0.18.1", optional = true }
 
 libc = { version = "0.2", optional = true }
 

dolby_vision/src/rpu/extension_metadata/blocks/level5.rs

@@ -12,7 +12,7 @@ const MAX_RESOLUTION_13_BITS: u16 = 8191;
 
 /// Active area of the picture (letterbox, aspect ratio)
 #[repr(C)]
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
 pub struct ExtMetadataBlockLevel5 {
     pub active_area_left_offset: u16,

src/commands/export.rs

@@ -1,6 +1,11 @@
-use clap::{Args, ValueHint};
 use std::path::PathBuf;
 
+use clap::{
+    builder::{EnumValueParser, PossibleValue, TypedValueParser},
+    Args, ValueEnum, ValueHint,
+};
+use clap_lex::OsStrExt as _;
+
 #[derive(Args, Debug)]
 pub struct ExportArgs {
     #[arg(
@@ -23,12 +28,78 @@ pub struct ExportArgs {
     )]
     pub input_pos: Option<PathBuf>,
 
+    #[arg(
+        id = "data",
+        help = "List of key-value export parameters formatted as `key=output`, where `output` is an output file path.\nSupports multiple occurences prefixed by --data or delimited by ','",
+        long,
+        short = 'd',
+        conflicts_with = "output",
+        value_parser = ExportOptionParser,
+        value_delimiter = ','
+    )]
+    pub data: Vec<(ExportData, Option<PathBuf>)>,
+
+    // FIXME: export single output deprecation
     #[arg(
         id = "output",
+        help = "Output JSON file name. Deprecated, replaced by `--data all=output`",
         long,
         short = 'o',
-        help = "Output JSON file name",
+        conflicts_with = "data",
+        hide = true,
         value_hint = ValueHint::FilePath
     )]
     pub output: Option<PathBuf>,
 }
+
+#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ExportData {
+    /// Exports the list of RPUs as a JSON file
+    All,
+    /// Exports the frame indices at which `scene_refresh_flag` is set to 1
+    Scenes,
+    /// Exports the video's L5 metadata in the form of an `editor` config JSON
+    Level5,
+}
+
+impl ExportData {
+    pub fn default_output_file(&self) -> &'static str {
+        match self {
+            ExportData::All => "RPU_export.json",
+            ExportData::Scenes => "RPU_scenes.txt",
+            ExportData::Level5 => "RPU_L5_edit_config.json",
+        }
+    }
+}
+
+#[derive(Clone)]
+struct ExportOptionParser;
+impl TypedValueParser for ExportOptionParser {
+    type Value = (ExportData, Option<PathBuf>);
+
+    fn parse_ref(
+        &self,
+        cmd: &clap::Command,
+        arg: Option<&clap::Arg>,
+        value: &std::ffi::OsStr,
+    ) -> Result<Self::Value, clap::Error> {
+        let data_parser = EnumValueParser::<ExportData>::new();
+
+        if let Some((data_str, output_str)) = value.split_once("=") {
+            Ok((
+                data_parser.parse_ref(cmd, arg, data_str)?,
+                output_str.to_str().map(str::parse).and_then(Result::ok),
+            ))
+        } else {
+            Ok((data_parser.parse_ref(cmd, arg, value)?, None))
+        }
+    }
+
+    fn possible_values(&self) -> Option<Box<dyn Iterator<Item = PossibleValue> + '_>> {
+        Some(Box::new(
+            ExportData::value_variants()
+                .iter()
+                .filter_map(|v| v.to_possible_value()),
+        ))
+    }
+}

src/commands/mod.rs

@@ -16,7 +16,7 @@ mod plot;
 pub use convert::ConvertArgs;
 pub use demux::DemuxArgs;
 pub use editor::EditorArgs;
-pub use export::ExportArgs;
+pub use export::{ExportArgs, ExportData};
 pub use extract_rpu::ExtractRpuArgs;
 pub use generate::{ArgHdr10PlusPeakBrightnessSource, GenerateArgs};
 pub use info::InfoArgs;

src/dovi/exporter.rs

@@ -1,43 +1,45 @@
+use std::borrow::Cow;
 use std::fs::File;
 use std::io::{stdout, BufWriter, Write};
+use std::ops::Range;
 use std::path::PathBuf;
 
 use anyhow::Result;
+use dolby_vision::rpu::extension_metadata::blocks::{ExtMetadataBlock, ExtMetadataBlockLevel5};
+use itertools::Itertools;
 use serde::ser::SerializeSeq;
 use serde::Serializer;
 
 use dolby_vision::rpu::utils::parse_rpu_file;
+use serde_json::json;
 
-use crate::commands::ExportArgs;
+use crate::commands::{ExportArgs, ExportData};
 use crate::dovi::input_from_either;
 
 use super::DoviRpu;
 
 pub struct Exporter {
     input: PathBuf,
-    output: PathBuf,
+    data: Vec<(ExportData, Option<PathBuf>)>,
 }
 
 impl Exporter {
     pub fn export(args: ExportArgs) -> Result<()> {
         let ExportArgs {
             input,
             input_pos,
+            data,
             output,
         } = args;
 
         let input = input_from_either("editor", input, input_pos)?;
+        let mut exporter = Exporter { input, data };
 
-        let out_path = if let Some(out_path) = output {
-            out_path
-        } else {
-            PathBuf::from("RPU_export.json".to_string())
-        };
+        if exporter.data.is_empty() {
+            exporter.data.push((ExportData::All, output));
+        }
 
-        let exporter = Exporter {
-            input,
-            output: out_path,
-        };
+        exporter.data.dedup_by_key(|(k, _)| *k);
 
         println!("Parsing RPU file...");
         stdout().flush().ok();
@@ -51,20 +53,131 @@ impl Exporter {
     }
 
     fn execute(&self, rpus: &[DoviRpu]) -> Result<()> {
-        println!("Exporting metadata...");
+        for (data, maybe_output) in &self.data {
+            let out_path = if let Some(out_path) = maybe_output {
+                Cow::Borrowed(out_path)
+            } else {
+                Cow::Owned(PathBuf::from(data.default_output_file()))
+            };
+
+            let writer_buf_len = if matches!(data, ExportData::All) {
+                100_000
+            } else {
+                1000
+            };
+            let mut writer = BufWriter::with_capacity(
+                writer_buf_len,
+                File::create(out_path.as_path()).expect("Can't create file"),
+            );
+
+            match data {
+                ExportData::All => {
+                    println!("Exporting serialized RPU list...");
+
+                    let mut ser = serde_json::Serializer::new(&mut writer);
+                    let mut seq = ser.serialize_seq(Some(rpus.len()))?;
+
+                    for rpu in rpus {
+                        seq.serialize_element(&rpu)?;
+                    }
+                    seq.end()?;
+                }
+                ExportData::Scenes => {
+                    println!("Exporting scenes list...");
+
+                    let scene_refresh_indices = rpus
+                        .iter()
+                        .enumerate()
+                        .filter(|(_, rpu)| {
+                            rpu.vdr_dm_data
+                                .as_ref()
+                                .is_some_and(|vdr| vdr.scene_refresh_flag == 1)
+                        })
+                        .map(|e| e.0);
+                    for i in scene_refresh_indices {
+                        writeln!(&mut writer, "{i}")?;
+                    }
+                }
+                ExportData::Level5 => {
+                    self.export_level5_config(rpus, &mut writer)?;
+                }
+            }
+
+            writer.flush()?;
+        }
 
-        let writer = BufWriter::with_capacity(
-            100_000,
-            File::create(&self.output).expect("Can't create file"),
-        );
+        Ok(())
+    }
 
-        let mut ser = serde_json::Serializer::new(writer);
-        let mut seq = ser.serialize_seq(Some(rpus.len()))?;
+    fn export_level5_config<W: Write>(&self, rpus: &[DoviRpu], writer: &mut W) -> Result<()> {
+        println!("Exporting L5 metadata config...");
+
+        let default_l5 = ExtMetadataBlockLevel5::default();
+
+        let l5_groups = rpus.iter().enumerate().group_by(|(_, rpu)| {
+            rpu.vdr_dm_data
+                .as_ref()
+                .and_then(|vdr| {
+                    vdr.get_block(5).and_then(|b| match b {
+                        ExtMetadataBlock::Level5(b) => Some(b),
+                        _ => None,
+                    })
+                })
+                .unwrap_or(&default_l5)
+        });
+        let l5_indices = l5_groups
+            .into_iter()
+            .map(|(k, group)| (k, group.take(1).map(|(i, _)| i).next().unwrap()));
+
+        let mut l5_presets =
+            Vec::<&ExtMetadataBlockLevel5>::with_capacity(l5_indices.size_hint().0);
+        let mut l5_edits = Vec::<(Range<usize>, usize)>::new();
+
+        for (k, start_index) in l5_indices {
+            if !l5_presets.iter().any(|l5| *l5 == k) {
+                l5_presets.push(k);
+            }
+
+            if let Some(last_edit) = l5_edits.last_mut() {
+                last_edit.0.end = start_index - 1;
+            }
+
+            let preset_idx = l5_presets.iter().position(|l5| *l5 == k).unwrap();
+            l5_edits.push((start_index..start_index, preset_idx));
+        }
 
-        for rpu in rpus {
-            seq.serialize_element(&rpu)?;
+        // Set last edit end index
+        if let Some(last_edit) = l5_edits.last_mut() {
+            last_edit.0.end = rpus.len() - 1;
         }
-        seq.end()?;
+
+        let l5_presets = l5_presets
+            .iter()
+            .enumerate()
+            .map(|(id, l5)| {
+                json!({
+                    "id": id,
+                    "left": l5.active_area_left_offset,
+                    "right": l5.active_area_right_offset,
+                    "top": l5.active_area_top_offset,
+                    "bottom": l5.active_area_bottom_offset
+                })
+            })
+            .collect::<Vec<_>>();
+        let l5_edits = l5_edits.iter().map(|(edit_range, id)| {
+            (
+                format!("{}-{}", edit_range.start, edit_range.end),
+                json!(id),
+            )
+        });
+        let l5_edits = serde_json::Value::Object(l5_edits.collect());
+
+        let edit_config = json!({
+            "crop": true,
+            "presets": l5_presets,
+            "edits": l5_edits,
+        });
+        serde_json::to_writer_pretty(writer, &edit_config)?;
 
         Ok(())
     }