Skip to content

Commit

Permalink
Add support for passing variables to the vspipe python environment (m…
Browse files Browse the repository at this point in the history 
…aster-of-zen#858)

* feat: new vspipe_args argument

makes it possible to pass variables to the vapoursynth python environment

* refactor: a more intrusive implementation

* refactor: named args for Vapoursynth and Video variants

* fix: use vspipe_args from Input in scene_detect.rs and chunk.rs

* fix: append vspipe_args into the main command as well

* fix: cargo fmt lints
  • Loading branch information
@Vernoxvernax
Vernoxvernax authored Jul 6, 2024
1 parent e9d953d commit 01c5781
Show file tree
Hide file tree
Showing 10 changed files with 211 additions and 91 deletions.
12 changes: 9 additions & 3 deletions av1an-core/src/chunk.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,9 @@ mod tests {
let ch = Chunk {
temp: "none".to_owned(),
index: 1,
input: Input::Video("test.mkv".into()),
input: Input::Video {
path: "test.mkv".into(),
},
source_cmd: vec!["".into()],
output_ext: "ivf".to_owned(),
start_frame: 0,
Expand All @@ -120,7 +122,9 @@ mod tests {
let ch = Chunk {
temp: "none".to_owned(),
index: 10000,
input: Input::Video("test.mkv".into()),
input: Input::Video {
path: "test.mkv".into(),
},
source_cmd: vec!["".into()],
output_ext: "ivf".to_owned(),
start_frame: 0,
Expand All @@ -141,7 +145,9 @@ mod tests {
let ch = Chunk {
temp: "d".to_owned(),
index: 1,
input: Input::Video("test.mkv".into()),
input: Input::Video {
path: "test.mkv".into(),
},
source_cmd: vec!["".into()],
output_ext: "ivf".to_owned(),
start_frame: 0,
Expand Down
52 changes: 34 additions & 18 deletions av1an-core/src/context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -151,23 +151,28 @@ impl Av1anContext {
&& !self.args.resume
{
self.vs_script = Some(match &self.args.input {
Input::VapourSynth(path) => path.clone(),
Input::Video(path) => create_vs_file(&self.args.temp, path, self.args.chunk_method)?,
Input::VapourSynth { path, .. } => path.clone(),
Input::Video{ path } => create_vs_file(&self.args.temp, path, self.args.chunk_method)?,
});

let vs_script = self.vs_script.clone().unwrap();
let vspipe_args = self.args.input.as_vspipe_args_vec()?;
Some({
thread::spawn(move || {
Command::new("vspipe")
.arg("-i")
.arg(vs_script)
.args(["-i", "-"])
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.unwrap()
.wait()
.unwrap()
let mut command = Command::new("vspipe");
command.arg("-i")
.arg(vs_script)
.args(["-i", "-"])
.stdout(Stdio::piped())
.stderr(Stdio::piped());
// Append vspipe arguments to the environment if there are any
for arg in vspipe_args {
command.args(["-a", &arg]);
}
command.spawn()
.unwrap()
.wait()
.unwrap()
})
})
} else {
Expand Down Expand Up @@ -455,7 +460,11 @@ impl Av1anContext {
let (source_pipe_stderr, ffmpeg_pipe_stderr, enc_output, enc_stderr, frame) =
rt.block_on(async {
let mut source_pipe = if let [source, args @ ..] = &*chunk.source_cmd {
tokio::process::Command::new(source)
let mut command = tokio::process::Command::new(source);
for arg in chunk.input.as_vspipe_args_vec().unwrap() {
command.args(["-a", &arg]);
}
command
.args(args)
.stdout(Stdio::piped())
.stderr(Stdio::piped())
Expand Down Expand Up @@ -666,7 +675,7 @@ impl Av1anContext {

fn create_encoding_queue(&mut self, scenes: &[Scene]) -> anyhow::Result<Vec<Chunk>> {
let mut chunks = match &self.args.input {
Input::Video(_) => match self.args.chunk_method {
Input::Video { .. } => match self.args.chunk_method {
ChunkMethod::FFMS2
| ChunkMethod::LSMASH
| ChunkMethod::DGDECNV
Expand All @@ -678,7 +687,7 @@ impl Av1anContext {
ChunkMethod::Select => self.create_video_queue_select(scenes),
ChunkMethod::Segment => self.create_video_queue_segment(scenes)?,
},
Input::VapourSynth(vs_script) => self.create_video_queue_vs(scenes, vs_script.as_path()),
Input::VapourSynth { path, .. } => self.create_video_queue_vs(scenes, path.as_path()),
};

match self.args.chunk_order {
Expand Down Expand Up @@ -878,7 +887,9 @@ impl Av1anContext {
let mut chunk = Chunk {
temp: self.args.temp.clone(),
index,
input: Input::Video(src_path.to_path_buf()),
input: Input::Video {
path: src_path.to_path_buf(),
},
source_cmd: ffmpeg_gen_cmd,
output_ext: output_ext.to_owned(),
start_frame,
Expand Down Expand Up @@ -931,7 +942,10 @@ impl Av1anContext {
let mut chunk = Chunk {
temp: self.args.temp.clone(),
index,
input: Input::VapourSynth(vs_script.to_path_buf()),
input: Input::VapourSynth {
path: vs_script.to_path_buf(),
vspipe_args: self.args.input.as_vspipe_args_vec()?,
},
source_cmd: vspipe_cmd_gen,
output_ext: output_ext.to_owned(),
start_frame: scene.start_frame,
Expand Down Expand Up @@ -1130,7 +1144,9 @@ impl Av1anContext {

let mut chunk = Chunk {
temp: self.args.temp.clone(),
input: Input::Video(PathBuf::from(file)),
input: Input::Video {
path: PathBuf::from(file),
},
source_cmd: ffmpeg_gen_cmd,
output_ext: output_ext.to_owned(),
index,
Expand Down
104 changes: 71 additions & 33 deletions av1an-core/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,9 @@ use std::thread::available_parallelism;
use std::time::Instant;

use ::ffmpeg::color::TransferCharacteristic;
use anyhow::Context;
use ::vapoursynth::api::API;
use ::vapoursynth::map::OwnedMap;
use anyhow::{bail, Context};
use av1_grain::TransferFunction;
use chunk::Chunk;
use dashmap::DashMap;
Expand Down Expand Up @@ -63,16 +65,21 @@ pub mod vmaf;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Input {
VapourSynth(PathBuf),
Video(PathBuf),
VapourSynth {
path: PathBuf,
vspipe_args: Vec<String>,
},
Video {
path: PathBuf,
},
}

impl Input {
/// Returns a reference to the inner path, panicking if the input is not an `Input::Video`.
pub fn as_video_path(&self) -> &Path {
match &self {
Input::Video(path) => path.as_ref(),
Input::VapourSynth(_) => {
Input::Video { path } => path.as_ref(),
Input::VapourSynth { .. } => {
panic!("called `Input::as_video_path()` on an `Input::VapourSynth` variant")
}
}
Expand All @@ -81,8 +88,8 @@ impl Input {
/// Returns a reference to the inner path, panicking if the input is not an `Input::VapourSynth`.
pub fn as_vapoursynth_path(&self) -> &Path {
match &self {
Input::VapourSynth(path) => path.as_ref(),
Input::Video(_) => {
Input::VapourSynth { path, .. } => path.as_ref(),
Input::Video { .. } => {
panic!("called `Input::as_vapoursynth_path()` on an `Input::Video` variant")
}
}
Expand All @@ -96,62 +103,66 @@ impl Input {
/// input type!
pub fn as_path(&self) -> &Path {
match &self {
Input::Video(path) | Input::VapourSynth(path) => path.as_ref(),
Input::Video { path } | Input::VapourSynth { path, .. } => path.as_ref(),
}
}

pub const fn is_video(&self) -> bool {
matches!(&self, Input::Video(_))
matches!(&self, Input::Video { .. })
}

pub const fn is_vapoursynth(&self) -> bool {
matches!(&self, Input::VapourSynth(_))
matches!(&self, Input::VapourSynth { .. })
}

pub fn frames(&self) -> anyhow::Result<usize> {
const FAIL_MSG: &str = "Failed to get number of frames for input video";
Ok(match &self {
Input::Video(path) => {
Input::Video { path } => {
ffmpeg::num_frames(path.as_path()).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
}
Input::VapourSynth(path) => {
vapoursynth::num_frames(path.as_path()).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
Input::VapourSynth { path, .. } => {
vapoursynth::num_frames(path.as_path(), self.as_vspipe_args_map()?)
.map_err(|_| anyhow::anyhow!(FAIL_MSG))?
}
})
}

pub fn frame_rate(&self) -> anyhow::Result<f64> {
const FAIL_MSG: &str = "Failed to get frame rate for input video";
Ok(match &self {
Input::Video(path) => {
Input::Video { path } => {
crate::ffmpeg::frame_rate(path.as_path()).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
}
Input::VapourSynth(path) => {
vapoursynth::frame_rate(path.as_path()).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
Input::VapourSynth { path, .. } => {
vapoursynth::frame_rate(path.as_path(), self.as_vspipe_args_map()?)
.map_err(|_| anyhow::anyhow!(FAIL_MSG))?
}
})
}

pub fn resolution(&self) -> anyhow::Result<(u32, u32)> {
const FAIL_MSG: &str = "Failed to get resolution for input video";
Ok(match self {
Input::VapourSynth(video) => {
crate::vapoursynth::resolution(video).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
Input::VapourSynth { path, .. } => {
crate::vapoursynth::resolution(path, self.as_vspipe_args_map()?)
.map_err(|_| anyhow::anyhow!(FAIL_MSG))?
}
Input::Video(video) => {
crate::ffmpeg::resolution(video).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
Input::Video { path } => {
crate::ffmpeg::resolution(path).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
}
})
}

pub fn pixel_format(&self) -> anyhow::Result<String> {
const FAIL_MSG: &str = "Failed to get resolution for input video";
Ok(match self {
Input::VapourSynth(video) => {
crate::vapoursynth::pixel_format(video).map_err(|_| anyhow::anyhow!(FAIL_MSG))?
Input::VapourSynth { path, .. } => {
crate::vapoursynth::pixel_format(path, self.as_vspipe_args_map()?)
.map_err(|_| anyhow::anyhow!(FAIL_MSG))?
}
Input::Video(video) => {
let fmt = crate::ffmpeg::get_pixel_format(video).map_err(|_| anyhow::anyhow!(FAIL_MSG))?;
Input::Video { path } => {
let fmt = crate::ffmpeg::get_pixel_format(path).map_err(|_| anyhow::anyhow!(FAIL_MSG))?;
format!("{fmt:?}")
}
})
Expand All @@ -160,16 +171,16 @@ impl Input {
fn transfer_function(&self) -> anyhow::Result<TransferFunction> {
const FAIL_MSG: &str = "Failed to get transfer characteristics for input video";
Ok(match self {
Input::VapourSynth(video) => {
match crate::vapoursynth::transfer_characteristics(video)
Input::VapourSynth { path, .. } => {
match crate::vapoursynth::transfer_characteristics(path, self.as_vspipe_args_map()?)
.map_err(|_| anyhow::anyhow!(FAIL_MSG))?
{
16 => TransferFunction::SMPTE2084,
_ => TransferFunction::BT1886,
}
}
Input::Video(video) => {
match crate::ffmpeg::transfer_characteristics(video)
Input::Video { path } => {
match crate::ffmpeg::transfer_characteristics(path)
.map_err(|_| anyhow::anyhow!(FAIL_MSG))?
{
TransferCharacteristic::SMPTE2084 => TransferFunction::SMPTE2084,
Expand Down Expand Up @@ -220,19 +231,46 @@ impl Input {
_ => (1, 1),
}
}

/// Returns the vector of arguments passed to the vspipe python environment
/// If the input is not a vapoursynth script, the vector will be empty.
pub fn as_vspipe_args_vec(&self) -> Result<Vec<String>, anyhow::Error> {
match self {
Input::VapourSynth { vspipe_args, .. } => Ok(vspipe_args.to_owned()),
Input::Video { .. } => Ok(vec![]),
}
}

/// Creates and returns an OwnedMap of the arguments passed to the vspipe python environment
/// If the input is not a vapoursynth script, the map will be empty.
pub fn as_vspipe_args_map(&self) -> Result<OwnedMap<'static>, anyhow::Error> {
let mut args_map = OwnedMap::new(API::get().unwrap());

for arg in self.as_vspipe_args_vec()? {
let split: Vec<&str> = arg.split_terminator('=').collect();
if args_map.set_data(split[0], split[1].as_bytes()).is_err() {
bail!("Failed to split vspipe arguments");
};
}

Ok(args_map)
}
}

impl<P: AsRef<Path> + Into<PathBuf>> From<P> for Input {
impl<P: AsRef<Path> + Into<PathBuf>> From<(P, Vec<String>)> for Input {
#[allow(clippy::option_if_let_else)]
fn from(path: P) -> Self {
fn from((path, vspipe_args): (P, Vec<String>)) -> Self {
if let Some(ext) = path.as_ref().extension() {
if ext == "py" || ext == "vpy" {
Self::VapourSynth(path.into())
Self::VapourSynth {
path: path.into(),
vspipe_args,
}
} else {
Self::Video(path.into())
Self::Video { path: path.into() }
}
} else {
Self::Video(path.into())
Self::Video { path: path.into() }
}
}
}
Expand Down
22 changes: 13 additions & 9 deletions av1an-core/src/scene_detect.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -229,21 +229,25 @@ fn build_decoder(
};

let decoder = match input {
Input::VapourSynth(path) => {
bit_depth = crate::vapoursynth::bit_depth(path.as_ref())?;
Input::VapourSynth { path, .. } => {
bit_depth = crate::vapoursynth::bit_depth(path.as_ref(), input.as_vspipe_args_map()?)?;
let vspipe_args = input.as_vspipe_args_vec()?;

if !filters.is_empty() {
let vspipe = Command::new("vspipe")
if !filters.is_empty() || !vspipe_args.is_empty() {
let mut command = Command::new("vspipe");
command
.arg("-c")
.arg("y4m")
.arg(path)
.arg("-")
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()?
.stdout
.unwrap();
.stderr(Stdio::null());
// Append vspipe python arguments to the environment if there are any
for arg in vspipe_args {
command.args(["-a", &arg]);
}
let vspipe = command.spawn()?.stdout.unwrap();
Decoder::Y4m(y4m::Decoder::new(
Command::new("ffmpeg")
.stdin(vspipe)
Expand All @@ -260,7 +264,7 @@ fn build_decoder(
Decoder::Vapoursynth(VapoursynthDecoder::new(path.as_ref())?)
}
}
Input::Video(path) => {
Input::Video { path } => {
let input_pix_format = crate::ffmpeg::get_pixel_format(path.as_ref())
.unwrap_or_else(|e| panic!("FFmpeg failed to get pixel format for input video: {e:?}"));
bit_depth = encoder.get_format_bit_depth(sc_pix_format.unwrap_or(input_pix_format))?;
Expand Down
4 changes: 3 additions & 1 deletion av1an-core/src/scenes.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -278,7 +278,9 @@ fn get_test_args() -> Av1anContext {
input_pix_format: InputPixelFormat::FFmpeg {
format: Pixel::YUV420P10LE,
},
input: Input::Video(PathBuf::new()),
input: Input::Video {
path: PathBuf::new(),
},
output_pix_format: PixelFormat {
format: Pixel::YUV420P10LE,
bit_depth: 10,
Expand Down
Loading

0 comments on commit 01c5781

Please sign in to comment.