From e8ae09c73124fd711ae8d635d1f240af93d796ae Mon Sep 17 00:00:00 2001 From: quietvoid <39477805+quietvoid@users.noreply.github.com> Date: Sat, 7 Oct 2023 09:14:30 -0400 Subject: [PATCH] Add remove subcommand Fixes #232 --- README.md | 15 ++++++++- src/commands/mod.rs | 7 +++- src/commands/remove.rs | 33 ++++++++++++++++++ src/dovi/mod.rs | 1 + src/dovi/remover.rs | 59 ++++++++++++++++++++++++++++++++ src/main.rs | 26 ++++++++------- tests/hevc/mod.rs | 1 + tests/hevc/remove.rs | 76 ++++++++++++++++++++++++++++++++++++++++++ 8 files changed, 204 insertions(+), 14 deletions(-) create mode 100644 src/commands/remove.rs create mode 100644 src/dovi/remover.rs create mode 100644 tests/hevc/remove.rs diff --git a/README.md b/README.md index 596af67..b7d1efd 100644 --- a/README.md +++ b/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. diff --git a/src/commands/mod.rs b/src/commands/mod.rs index 772ab7e..d84387f 100644 --- a/src/commands/mod.rs +++ b/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)] diff --git a/src/commands/remove.rs b/src/commands/remove.rs new file mode 100644 index 0000000..a89dadb --- /dev/null +++ b/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, + + #[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, + + #[arg( + long, + short = 'o', + help = "Base layer output file location", + value_hint = ValueHint::FilePath + )] + pub output: Option, +} diff --git a/src/dovi/mod.rs b/src/dovi/mod.rs index 9638a39..8e8254e 100644 --- a/src/dovi/mod.rs +++ b/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; diff --git a/src/dovi/remover.rs b/src/dovi/remover.rs new file mode 100644 index 0000000..5211d7c --- /dev/null +++ b/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 { + 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) + } +} diff --git a/src/main.rs b/src/main.rs index cff7812..19f08bb 100644 --- a/src/main.rs +++ b/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), } } diff --git a/tests/hevc/mod.rs b/tests/hevc/mod.rs index 62135ea..2762b9f 100644 --- a/tests/hevc/mod.rs +++ b/tests/hevc/mod.rs @@ -3,3 +3,4 @@ mod demux; mod extract_rpu; mod inject_rpu; mod mux; +mod remove; diff --git a/tests/hevc/remove.rs b/tests/hevc/remove.rs new file mode 100644 index 0000000..efb1b68 --- /dev/null +++ b/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(()) +}