Skip to content

Add support for passing variables to the vspipe python environment #858

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.

Sign up for GitHub

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

Merged
merged 6 commits into from
Jul 6, 2024
Merged

Add support for passing variables to the vspipe python environment #858

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
8cf54bc
feat: new vspipe_args argument
Vernoxvernax Jul 4, 2024
7302809
refactor: a more intrusive implementation
Vernoxvernax Jul 5, 2024
be91abd
refactor: named args for Vapoursynth and Video variants
Vernoxvernax Jul 5, 2024
e0f1e55
fix: use vspipe_args from Input in scene_detect.rs and chunk.rs
Vernoxvernax Jul 5, 2024
5828a44
fix: append vspipe_args into the main command as well
Vernoxvernax Jul 5, 2024
ab3295c
fix: cargo fmt lints
Vernoxvernax Jul 6, 2024
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
12 changes: 9 additions & 3 deletions av1an-core/src/chunk.rs
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading