Support zerocopy VS scene detection

master-of-zen · Jul 13, 2022 · 09c00d2 · 09c00d2
1 parent a63d14a
commit 09c00d2
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 96 deletions.
diff --git a/av1an-core/src/lib.rs b/av1an-core/src/lib.rs
@@ -72,6 +72,7 @@ pub enum Input {
 
 impl Input {
   /// Returns a reference to the inner path, panicking if the input is not an `Input::Video`.
+  #[track_caller]
   pub fn as_video_path(&self) -> &Path {
     match &self {
       Input::Video(path) => path.as_ref(),
@@ -82,6 +83,7 @@ impl Input {
   }
 
   /// Returns a reference to the inner path, panicking if the input is not an `Input::VapourSynth`.
+  #[track_caller]
   pub fn as_vapoursynth_path(&self) -> &Path {
     match &self {
       Input::VapourSynth(path) => path.as_ref(),

diff --git a/av1an-core/src/scene_detect.rs b/av1an-core/src/scene_detect.rs
@@ -1,17 +1,16 @@
-use std::io::Read;
-use std::process::{Command, Stdio};
 use std::thread;
 
 use ansi_term::Style;
 use anyhow::bail;
-use av_metrics_decoders::{Decoder2, FfmpegDecoder};
+use av_metrics_decoders::{Decoder2, FfmpegDecoder, VapoursynthDecoder};
 use av_scenechange::{detect_scene_changes, DetectionOptions, SceneDetectionSpeed};
 use ffmpeg::format::Pixel;
 use itertools::Itertools;
-use smallvec::{smallvec, SmallVec};
+
+use vapoursynth::prelude::*;
 
 use crate::scenes::Scene;
-use crate::{into_smallvec, progress_bar, Encoder, Input, ScenecutMethod, Verbosity};
+use crate::{progress_bar, Encoder, Input, ScenecutMethod, Verbosity};
 
 pub fn av_scenechange_detect(
   input: &Input,
@@ -43,27 +42,63 @@ pub fn av_scenechange_detect(
     frames
   });
 
-  let mut ctx = FfmpegDecoder::get_ctx(input.as_video_path()).unwrap();
+  let mut ff_ctx;
+  let mut ff_decoder;
 
-  let mut decoder = FfmpegDecoder::new(&mut ctx).unwrap();
+  let mut vs_ctx;
+  let mut vs_decoder;
 
-  let scenes = scene_detect(
-    &mut decoder,
-    encoder,
-    total_frames,
-    if verbosity == Verbosity::Quiet {
-      None
-    } else {
-      Some(&|frames, _| {
-        progress_bar::set_pos(frames as u64);
-      })
-    },
-    min_scene_len,
-    sc_pix_format,
-    sc_method,
-    sc_downscale_height,
-    zones,
-  )?;
+  let scenes = match input {
+    Input::Video(path) => {
+      ff_ctx = FfmpegDecoder::get_ctx(path).unwrap();
+      ff_decoder = FfmpegDecoder::new(&mut ff_ctx).unwrap();
+
+      scene_detect(
+        &mut ff_decoder,
+        encoder,
+        total_frames,
+        if verbosity == Verbosity::Quiet {
+          None
+        } else {
+          Some(&|frames, _| {
+            progress_bar::set_pos(frames as u64);
+          })
+        },
+        min_scene_len,
+        sc_pix_format,
+        sc_method,
+        sc_downscale_height,
+        zones,
+      )?
+    }
+    Input::VapourSynth(path) => {
+      // TODO make helper function for creating VS Environment
+      vs_ctx = Environment::new().unwrap();
+
+      // Evaluate the script.
+      vs_ctx.eval_file(path, EvalFlags::SetWorkingDir).unwrap();
+
+      vs_decoder = VapoursynthDecoder::new(&vs_ctx).unwrap();
+
+      scene_detect(
+        &mut vs_decoder,
+        encoder,
+        total_frames,
+        if verbosity == Verbosity::Quiet {
+          None
+        } else {
+          Some(&|frames, _| {
+            progress_bar::set_pos(frames as u64);
+          })
+        },
+        min_scene_len,
+        sc_pix_format,
+        sc_method,
+        sc_downscale_height,
+        zones,
+      )?
+    }
+  };
 
   let frames = frame_thread.join().unwrap();
 
@@ -197,75 +232,3 @@ pub fn scene_detect<F, D: Decoder2<F>>(
   }
   Ok(scenes)
 }
-
-#[allow(unused)]
-fn build_decoder(
-  input: &Input,
-  encoder: Encoder,
-  sc_pix_format: Option<Pixel>,
-  sc_downscale_height: Option<usize>,
-) -> anyhow::Result<(y4m::Decoder<impl Read>, usize)> {
-  let bit_depth;
-  let filters: SmallVec<[String; 4]> = match (sc_downscale_height, sc_pix_format) {
-    (Some(sdh), Some(spf)) => into_smallvec![
-      "-vf",
-      format!(
-        "format={},scale=-2:'min({},ih)'",
-        spf.descriptor().unwrap().name(),
-        sdh
-      )
-    ],
-    (Some(sdh), None) => into_smallvec!["-vf", format!("scale=-2:'min({},ih)'", sdh)],
-    (None, Some(spf)) => into_smallvec!["-pix_fmt", spf.descriptor().unwrap().name()],
-    (None, None) => smallvec![],
-  };
-
-  let decoder = y4m::Decoder::new(match input {
-    Input::VapourSynth(path) => {
-      bit_depth = crate::vapoursynth::bit_depth(path.as_ref())?;
-      let vspipe = Command::new("vspipe")
-        .arg("-y")
-        .arg(path)
-        .arg("-")
-        .stdin(Stdio::null())
-        .stdout(Stdio::piped())
-        .stderr(Stdio::null())
-        .spawn()?
-        .stdout
-        .unwrap();
-
-      if !filters.is_empty() {
-        Command::new("ffmpeg")
-          .stdin(vspipe)
-          .args(["-i", "pipe:", "-f", "yuv4mpegpipe", "-strict", "-1"])
-          .args(filters)
-          .arg("-")
-          .stdout(Stdio::piped())
-          .stderr(Stdio::null())
-          .spawn()?
-          .stdout
-          .unwrap()
-      } else {
-        vspipe
-      }
-    }
-    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))?;
-      Command::new("ffmpeg")
-        .args(["-r", "1", "-i"])
-        .arg(path)
-        .args(filters.as_ref())
-        .args(["-f", "yuv4mpegpipe", "-strict", "-1", "-"])
-        .stdin(Stdio::null())
-        .stdout(Stdio::piped())
-        .stderr(Stdio::null())
-        .spawn()?
-        .stdout
-        .unwrap()
-    }
-  })?;
-
-  Ok((decoder, bit_depth))
-}