Add remove subcommand (#257)

Fixes #232

README.md

@@ -202,7 +202,7 @@ For working with an HEVC source file, there are multiple options that apply to m
 * ### **demux**
     Rust port of yusesope's python tool. Credits goes to them.  
     Demuxes single track dual layer Dolby Vision into Base layer and Enhancement layer files.  
-    Also can be used to remove the RPUs from an HEVC file.
+    The base layer file output is equivalent to using the `remove` subcommand.
 
     **Flags**:
     - `--el-only` Output the EL file only.
@@ -283,6 +283,19 @@ For working with an HEVC source file, there are multiple options that apply to m
     dovi_tool inject-rpu -i video.hevc --rpu-in RPU.bin -o injected_output.hevc
     ```
 
+ 
+* ### **remove**
+    Removes the enhancement layer and RPU data from the video.  
+    Outputs to a `BL.hevc` file by default.
+
+    **Examples**:
+    ```console
+    dovi_tool remove file.hevc
+    ```
+    ```console
+    ffmpeg -i input.mkv -c:v copy -vbsf hevc_mp4toannexb -f hevc - | dovi_tool remove -
+    ```
+
  
 
 Build artifacts can be found in the Github Actions.  

src/commands/mod.rs

@@ -12,6 +12,7 @@ mod info;
 mod inject_rpu;
 mod mux;
 mod plot;
+mod remove;
 
 pub use convert::ConvertArgs;
 pub use demux::DemuxArgs;
@@ -23,9 +24,10 @@ pub use info::InfoArgs;
 pub use inject_rpu::InjectRpuArgs;
 pub use mux::MuxArgs;
 pub use plot::PlotArgs;
+pub use remove::RemoveArgs;
 
 #[derive(Parser, Debug)]
-pub enum Command {
+pub enum Commands {
     #[command(about = "Converts RPU within a single layer HEVC file")]
     Convert(ConvertArgs),
 
@@ -57,6 +59,9 @@ pub enum Command {
 
     #[command(about = "Plot the L1 dynamic brightness metadata")]
     Plot(PlotArgs),
+
+    #[command(about = "Removes the enhancement layer and RPU data from the video")]
+    Remove(RemoveArgs),
 }
 
 #[derive(clap::ValueEnum, Debug, Copy, Clone)]

src/commands/remove.rs

@@ -0,0 +1,33 @@
+use clap::{Args, ValueHint};
+use std::path::PathBuf;
+
+#[derive(Args, Debug)]
+pub struct RemoveArgs {
+    #[arg(
+        id = "input",
+        help = "Sets the input HEVC file to use, or piped with -",
+        long,
+        short = 'i',
+        conflicts_with = "input_pos",
+        required_unless_present = "input_pos",
+        value_hint = ValueHint::FilePath,
+    )]
+    pub input: Option<PathBuf>,
+
+    #[arg(
+        id = "input_pos",
+        help = "Sets the input HEVC file to use, or piped with - (positional)",
+        conflicts_with = "input",
+        required_unless_present = "input",
+        value_hint = ValueHint::FilePath
+    )]
+    pub input_pos: Option<PathBuf>,
+
+    #[arg(
+        long,
+        short = 'o',
+        help = "Base layer output file location",
+        value_hint = ValueHint::FilePath
+    )]
+    pub output: Option<PathBuf>,
+}

src/dovi/mod.rs

@@ -20,6 +20,7 @@ pub mod exporter;
 pub mod generator;
 pub mod muxer;
 pub mod plotter;
+pub mod remover;
 pub mod rpu_extractor;
 pub mod rpu_info;
 pub mod rpu_injector;

src/dovi/remover.rs

@@ -0,0 +1,59 @@
+use anyhow::{bail, Result};
+use indicatif::ProgressBar;
+use std::path::PathBuf;
+
+use crate::commands::RemoveArgs;
+
+use super::{general_read_write, input_from_either, CliOptions, IoFormat};
+
+use general_read_write::{DoviProcessor, DoviWriter};
+
+pub struct Remover {
+    format: IoFormat,
+    input: PathBuf,
+    output: PathBuf,
+}
+
+impl Remover {
+    pub fn from_args(args: RemoveArgs) -> Result<Self> {
+        let RemoveArgs {
+            input,
+            input_pos,
+            output,
+        } = args;
+
+        let input = input_from_either("remove", input, input_pos)?;
+        let format = hevc_parser::io::format_from_path(&input)?;
+
+        let output = output.unwrap_or(PathBuf::from("BL.hevc"));
+
+        Ok(Self {
+            format,
+            input,
+            output,
+        })
+    }
+
+    pub fn remove(args: RemoveArgs, options: CliOptions) -> Result<()> {
+        let remover = Remover::from_args(args)?;
+        remover.process_input(options)
+    }
+
+    fn process_input(&self, options: CliOptions) -> Result<()> {
+        let pb = super::initialize_progress_bar(&self.format, &self.input)?;
+
+        match self.format {
+            IoFormat::Matroska => bail!("Remover: Matroska input is unsupported"),
+            _ => self.remove_from_raw_hevc(pb, options),
+        }
+    }
+
+    fn remove_from_raw_hevc(&self, pb: ProgressBar, options: CliOptions) -> Result<()> {
+        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);
+
+        dovi_processor.read_write_from_io(&self.format)
+    }
+}

src/main.rs

@@ -7,7 +7,7 @@ use clap::{Parser, ValueHint};
 mod tests;
 
 mod commands;
-use commands::{Command, ConversionModeCli};
+use commands::{Commands, ConversionModeCli};
 
 mod dovi;
 use dovi::{
@@ -18,6 +18,7 @@ use dovi::{
     generator::Generator,
     muxer::Muxer,
     plotter::Plotter,
+    remover::Remover,
     rpu_extractor::RpuExtractor,
     rpu_info::RpuInfo,
     rpu_injector::RpuInjector,
@@ -74,7 +75,7 @@ struct Opt {
     start_code: WriteStartCodePreset,
 
     #[command(subcommand)]
-    cmd: Command,
+    cmd: Commands,
 }
 
 fn main() -> Result<()> {
@@ -101,15 +102,16 @@ fn main() -> Result<()> {
     }
 
     match opt.cmd {
-        Command::Demux(args) => Demuxer::demux(args, cli_options),
-        Command::Editor(args) => Editor::edit(args),
-        Command::Convert(args) => Converter::convert(args, cli_options),
-        Command::ExtractRpu(args) => RpuExtractor::extract_rpu(args, cli_options),
-        Command::InjectRpu(args) => RpuInjector::inject_rpu(args, cli_options),
-        Command::Info(args) => RpuInfo::info(args),
-        Command::Generate(args) => Generator::generate(args),
-        Command::Export(args) => Exporter::export(args),
-        Command::Mux(args) => Muxer::mux_el(args, cli_options),
-        Command::Plot(args) => Plotter::plot(args),
+        Commands::Demux(args) => Demuxer::demux(args, cli_options),
+        Commands::Editor(args) => Editor::edit(args),
+        Commands::Convert(args) => Converter::convert(args, cli_options),
+        Commands::ExtractRpu(args) => RpuExtractor::extract_rpu(args, cli_options),
+        Commands::InjectRpu(args) => RpuInjector::inject_rpu(args, cli_options),
+        Commands::Info(args) => RpuInfo::info(args),
+        Commands::Generate(args) => Generator::generate(args),
+        Commands::Export(args) => Exporter::export(args),
+        Commands::Mux(args) => Muxer::mux_el(args, cli_options),
+        Commands::Plot(args) => Plotter::plot(args),
+        Commands::Remove(args) => Remover::remove(args, cli_options),
     }
 }

tests/hevc/mod.rs

@@ -3,3 +3,4 @@ mod demux;
 mod extract_rpu;
 mod inject_rpu;
 mod mux;
+mod remove;

tests/hevc/remove.rs

@@ -0,0 +1,76 @@
+use std::path::Path;
+
+use anyhow::Result;
+use assert_cmd::Command;
+use assert_fs::prelude::*;
+use predicates::prelude::*;
+
+const SUBCOMMAND: &str = "remove";
+
+#[test]
+fn help() -> Result<()> {
+    let mut cmd = Command::cargo_bin(env!("CARGO_PKG_NAME"))?;
+    let assert = cmd.arg(SUBCOMMAND).arg("--help").assert();
+
+    assert
+        .success()
+        .stderr(predicate::str::is_empty())
+        .stdout(predicate::str::contains(
+            "dovi_tool remove [OPTIONS] [input_pos]",
+        ));
+    Ok(())
+}
+
+#[test]
+fn remove() -> Result<()> {
+    let mut cmd = Command::cargo_bin(env!("CARGO_PKG_NAME"))?;
+    let temp = assert_fs::TempDir::new().unwrap();
+
+    let input_file = Path::new("assets/hevc_tests/regular_start_code_4_muxed_el.hevc");
+    let expected_bl = Path::new("assets/hevc_tests/regular_bl_start_code_4.hevc");
+
+    let output_bl = temp.child("BL.hevc");
+
+    let assert = cmd
+        .arg(SUBCOMMAND)
+        .arg(input_file)
+        .arg("--output")
+        .arg(output_bl.as_ref())
+        .assert();
+
+    assert.success().stderr(predicate::str::is_empty());
+
+    output_bl
+        .assert(predicate::path::is_file())
+        .assert(predicate::path::eq_file(expected_bl));
+
+    Ok(())
+}
+
+#[test]
+fn annexb() -> Result<()> {
+    let mut cmd = Command::cargo_bin(env!("CARGO_PKG_NAME"))?;
+    let temp = assert_fs::TempDir::new().unwrap();
+
+    let input_file = Path::new("assets/hevc_tests/regular_start_code_4_muxed_el.hevc");
+    let expected_bl = Path::new("assets/hevc_tests/regular_demux_bl_annexb.hevc");
+
+    let output_bl = temp.child("BL.hevc");
+
+    let assert = cmd
+        .arg("--start-code")
+        .arg("annex-b")
+        .arg(SUBCOMMAND)
+        .arg(input_file)
+        .arg("--output")
+        .arg(output_bl.as_ref())
+        .assert();
+
+    assert.success().stderr(predicate::str::is_empty());
+
+    output_bl
+        .assert(predicate::path::is_file())
+        .assert(predicate::path::eq_file(expected_bl));
+
+    Ok(())
+}