Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add remove subcommand #257

Merged
merged 1 commit into from
Oct 7, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
7 changes: 6 additions & 1 deletion src/commands/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ mod info;
mod inject_rpu;
mod mux;
mod plot;
mod remove;

pub use convert::ConvertArgs;
pub use demux::DemuxArgs;
Expand All @@ -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),

Expand Down Expand Up @@ -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)]
Expand Down
33 changes: 33 additions & 0 deletions src/commands/remove.rs
Original file line number Diff line number Diff line change
@@ -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>,
}
1 change: 1 addition & 0 deletions src/dovi/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
59 changes: 59 additions & 0 deletions src/dovi/remover.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
26 changes: 14 additions & 12 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use clap::{Parser, ValueHint};
mod tests;

mod commands;
use commands::{Command, ConversionModeCli};
use commands::{Commands, ConversionModeCli};

mod dovi;
use dovi::{
Expand All @@ -18,6 +18,7 @@ use dovi::{
generator::Generator,
muxer::Muxer,
plotter::Plotter,
remover::Remover,
rpu_extractor::RpuExtractor,
rpu_info::RpuInfo,
rpu_injector::RpuInjector,
Expand Down Expand Up @@ -74,7 +75,7 @@ struct Opt {
start_code: WriteStartCodePreset,

#[command(subcommand)]
cmd: Command,
cmd: Commands,
}

fn main() -> Result<()> {
Expand All @@ -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),
}
}
1 change: 1 addition & 0 deletions tests/hevc/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,4 @@ mod demux;
mod extract_rpu;
mod inject_rpu;
mod mux;
mod remove;
76 changes: 76 additions & 0 deletions tests/hevc/remove.rs
Original file line number Diff line number Diff line change
@@ -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(())
}