From 574d0c15e3e674f37ef75aa27a4eb86ca573676a Mon Sep 17 00:00:00 2001 From: Josh Holmer Date: Tue, 21 May 2024 20:07:43 -0400 Subject: [PATCH] Add Vapoursynth decoding support to allow for avoiding piping (#168) --- .github/workflows/av-scenechange.yml | 141 ++++++++++---------- .gitignore | 1 + Cargo.lock | 29 ++++ Cargo.toml | 18 ++- src/decoder.rs | 40 ++++++ src/lib.rs | 33 +++-- src/main.rs | 12 +- src/vapoursynth.rs | 191 +++++++++++++++++++++++++++ 8 files changed, 373 insertions(+), 92 deletions(-) create mode 100644 src/decoder.rs create mode 100644 src/vapoursynth.rs diff --git a/.github/workflows/av-scenechange.yml b/.github/workflows/av-scenechange.yml index 3ac5c5e..4e741cb 100644 --- a/.github/workflows/av-scenechange.yml +++ b/.github/workflows/av-scenechange.yml @@ -9,41 +9,38 @@ on: - master jobs: - clippy-rustfmt: - runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 - - - name: Install stable - uses: dtolnay/rust-toolchain@stable - with: - components: clippy, rustfmt - - - name: Install nasm - env: - LINK: http://debian-archive.trafficmanager.net/debian/pool/main/n/nasm - NASM_VERSION: 2.15.05-1 - NASM_SHA256: >- - c860caec653b865d5b83359452d97b11f1b3ba5b18b07cac554cf72550b3bfc9 - run: | - curl -O "$LINK/nasm_${NASM_VERSION}_amd64.deb" - echo "$NASM_SHA256 nasm_${NASM_VERSION}_amd64.deb" | sha256sum --check - sudo dpkg -i "nasm_${NASM_VERSION}_amd64.deb" - - - name: Run rustfmt - run: | - cargo fmt -- --check --verbose - - - name: Run clippy - uses: clechasseur/rs-clippy-check@v3 - with: - args: -- -D warnings --verbose -A clippy::wrong-self-convention -A clippy::many_single_char_names -A clippy::upper-case-acronyms + - uses: actions/checkout@v4 + + - name: Install stable + uses: dtolnay/rust-toolchain@stable + with: + components: clippy, rustfmt + + - name: Install nasm + env: + LINK: http://debian-archive.trafficmanager.net/debian/pool/main/n/nasm + NASM_VERSION: 2.15.05-1 + NASM_SHA256: >- + c860caec653b865d5b83359452d97b11f1b3ba5b18b07cac554cf72550b3bfc9 + run: | + curl -O "$LINK/nasm_${NASM_VERSION}_amd64.deb" + echo "$NASM_SHA256 nasm_${NASM_VERSION}_amd64.deb" | sha256sum --check + sudo dpkg -i "nasm_${NASM_VERSION}_amd64.deb" + + - name: Run rustfmt + run: | + cargo fmt -- --check --verbose + + - name: Run clippy + uses: clechasseur/rs-clippy-check@v3 + with: + args: -- -D warnings --verbose -A clippy::wrong-self-convention -A clippy::many_single_char_names -A clippy::upper-case-acronyms build: - strategy: matrix: platform: [ubuntu-latest, windows-latest] @@ -51,46 +48,46 @@ jobs: runs-on: ${{ matrix.platform }} steps: - - uses: actions/checkout@v4 - - - name: Install stable - uses: dtolnay/rust-toolchain@stable - - - name: Install nasm for Ubuntu - if: matrix.platform == 'ubuntu-latest' - env: - LINK: http://debian-archive.trafficmanager.net/debian/pool/main/n/nasm - NASM_VERSION: 2.15.05-1 - NASM_SHA256: >- - c860caec653b865d5b83359452d97b11f1b3ba5b18b07cac554cf72550b3bfc9 - run: | - curl -O "$LINK/nasm_${NASM_VERSION}_amd64.deb" - echo "$NASM_SHA256 nasm_${NASM_VERSION}_amd64.deb" | sha256sum --check - sudo dpkg -i "nasm_${NASM_VERSION}_amd64.deb" - - - name: Install nasm for Windows - if: matrix.platform == 'windows-latest' - run: | - $NASM_VERSION="2.15.05" - $LINK="https://www.nasm.us/pub/nasm/releasebuilds/$NASM_VERSION/win64" - curl -LO "$LINK/nasm-$NASM_VERSION-win64.zip" - 7z e -y "nasm-$NASM_VERSION-win64.zip" -o"C:\nasm" - echo "C:\nasm" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append - - - name: Set MSVC x86_64 linker path - if: matrix.platform == 'windows-latest' - run: | - $LinkGlob = "VC\Tools\MSVC\*\bin\Hostx64\x64" - $env:PATH = "$env:PATH;${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer" - $LinkPath = vswhere -latest -products * -find "$LinkGlob" | - Select-Object -Last 1 - echo "$LinkPath" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append - - - name: Build - run: cargo build --all-features --tests --benches - - - name: Run tests - run: cargo test --all-features - - - name: Generate docs - run: cargo doc --all-features --no-deps + - uses: actions/checkout@v4 + + - name: Install stable + uses: dtolnay/rust-toolchain@stable + + - name: Install nasm for Ubuntu + if: matrix.platform == 'ubuntu-latest' + env: + LINK: http://debian-archive.trafficmanager.net/debian/pool/main/n/nasm + NASM_VERSION: 2.15.05-1 + NASM_SHA256: >- + c860caec653b865d5b83359452d97b11f1b3ba5b18b07cac554cf72550b3bfc9 + run: | + curl -O "$LINK/nasm_${NASM_VERSION}_amd64.deb" + echo "$NASM_SHA256 nasm_${NASM_VERSION}_amd64.deb" | sha256sum --check + sudo dpkg -i "nasm_${NASM_VERSION}_amd64.deb" + + - name: Install nasm for Windows + if: matrix.platform == 'windows-latest' + run: | + $NASM_VERSION="2.15.05" + $LINK="https://www.nasm.us/pub/nasm/releasebuilds/$NASM_VERSION/win64" + curl -LO "$LINK/nasm-$NASM_VERSION-win64.zip" + 7z e -y "nasm-$NASM_VERSION-win64.zip" -o"C:\nasm" + echo "C:\nasm" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append + + - name: Set MSVC x86_64 linker path + if: matrix.platform == 'windows-latest' + run: | + $LinkGlob = "VC\Tools\MSVC\*\bin\Hostx64\x64" + $env:PATH = "$env:PATH;${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer" + $LinkPath = vswhere -latest -products * -find "$LinkGlob" | + Select-Object -Last 1 + echo "$LinkPath" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append + + - name: Build + run: cargo build --features binary,devel,tracing,serialize --tests --benches + + - name: Run tests + run: cargo test --features binary,devel,tracing,serialize + + - name: Generate docs + run: cargo doc --features binary,devel,tracing,serialize --no-deps diff --git a/.gitignore b/.gitignore index ebdff5a..b7aeb6d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ /target **/*.rs.bk /.idea +/.vscode diff --git a/Cargo.lock b/Cargo.lock index 40e5e9f..86842c3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -100,6 +100,7 @@ dependencies = [ "tracing", "tracing-chrome", "tracing-subscriber", + "vapoursynth", "y4m", ] @@ -117,6 +118,12 @@ dependencies = [ "v_frame", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitstream-io" version = "2.2.0" @@ -917,6 +924,28 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" +[[package]] +name = "vapoursynth" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c7df702c65dec1cfa3b93f824a1e58d5b0fdb82ac8a722596f43d7214282f56" +dependencies = [ + "anyhow", + "bitflags", + "lazy_static", + "thiserror", + "vapoursynth-sys", +] + +[[package]] +name = "vapoursynth-sys" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fc59bbb7980ce21ece45bc5b3e316bf27ac164b7b1c273ce4846c29d0642a9c" +dependencies = [ + "cfg-if", +] + [[package]] name = "version-compare" version = "0.1.1" diff --git a/Cargo.toml b/Cargo.toml index 9fce64a..2e5593b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,9 +14,9 @@ clap = { version = "4.0.22", optional = true, features = ["derive"] } serde = { version = "1.0.123", optional = true, features = ["derive"] } serde_json = { version = "1.0.62", optional = true } rav1e = { version = "0.7.0", default-features = false, features = [ - "asm", - "scenechange", - "threading", + "asm", + "scenechange", + "threading", ] } log = { version = "0.4.14", optional = true } console = { version = "0.15", optional = true } @@ -25,6 +25,16 @@ tracing-subscriber = { version = "0.3.18", optional = true } tracing-chrome = { version = "0.7.1", optional = true } tracing = { version = "0.1.40", optional = true } +[dependencies.vapoursynth] +version = "0.4.0" +features = [ + "vsscript-functions", + "vapoursynth-functions", + "vapoursynth-api-32", + "vsscript-api-31", +] +optional = true + [features] default = ["binary"] binary = ["clap", "serialize"] @@ -34,7 +44,7 @@ tracing = [ "tracing-subscriber", "tracing-chrome", "dep:tracing", - "rav1e/tracing" + "rav1e/tracing", ] [[bin]] diff --git a/src/decoder.rs b/src/decoder.rs new file mode 100644 index 0000000..bee3de7 --- /dev/null +++ b/src/decoder.rs @@ -0,0 +1,40 @@ +use std::io::Read; + +use rav1e::{Frame, Pixel}; + +#[cfg(feature = "vapoursynth")] +use crate::vapoursynth::VapoursynthDecoder; +use crate::y4m::VideoDetails; + +pub enum Decoder { + Y4m(y4m::Decoder), + #[cfg(feature = "vapoursynth")] + Vapoursynth(VapoursynthDecoder), +} + +impl Decoder { + /// # Errors + /// + /// - If using a Vapoursynth script that contains an unsupported video format. + pub fn get_video_details(&self) -> anyhow::Result { + match self { + Decoder::Y4m(dec) => Ok(crate::y4m::get_video_details(dec)), + #[cfg(feature = "vapoursynth")] + Decoder::Vapoursynth(dec) => dec.get_video_details(), + } + } + + /// # Errors + /// + /// - If a frame cannot be read. + pub fn read_video_frame( + &mut self, + video_details: &VideoDetails, + ) -> anyhow::Result> { + match self { + Decoder::Y4m(dec) => crate::y4m::read_video_frame::(dec, video_details), + #[cfg(feature = "vapoursynth")] + Decoder::Vapoursynth(dec) => dec.read_video_frame::(video_details), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 666aacc..788fa1d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -89,6 +89,9 @@ #![warn(clippy::missing_errors_doc)] #![warn(clippy::missing_panics_doc)] +pub mod decoder; +#[cfg(feature = "vapoursynth")] +pub mod vapoursynth; mod y4m; use std::{ @@ -98,7 +101,8 @@ use std::{ time::Instant, }; -use ::y4m::Decoder; +pub use ::y4m::Decoder as Y4mDecoder; +use decoder::Decoder; pub use rav1e::scenechange::SceneChangeDetector; use rav1e::{ config::{CpuFeatureLevel, EncoderConfig}, @@ -152,11 +156,14 @@ pub struct DetectionResults { pub speed: f64, } +/// # Errors +/// +/// - If using a Vapoursynth script that contains an unsupported video format. pub fn new_detector( dec: &mut Decoder, opts: DetectionOptions, -) -> SceneChangeDetector { - let video_details = y4m::get_video_details(dec); +) -> anyhow::Result> { + let video_details = dec.get_video_details()?; let mut config = EncoderConfig::with_speed_preset(if opts.analysis_speed == SceneDetectionSpeed::Fast { 10 @@ -178,7 +185,7 @@ pub fn new_detector( config.speed_settings.transform.tx_domain_distortion = true; let sequence = Arc::new(Sequence::new(&config)); - SceneChangeDetector::new( + Ok(SceneChangeDetector::new( config, CpuFeatureLevel::default(), if opts.detect_flashes { @@ -187,7 +194,7 @@ pub fn new_detector( 1 }, sequence, - ) + )) } /// Runs through a y4m video clip, @@ -204,6 +211,10 @@ pub fn new_detector( /// analyzed, and the number of keyframes detected. This is generally useful /// for displaying progress, etc. /// +/// # Errors +/// +/// - If using a Vapoursynth script that contains an unsupported video format. +/// /// # Panics /// /// - If `opts.lookahead_distance` is 0. @@ -213,11 +224,11 @@ pub fn detect_scene_changes( opts: DetectionOptions, frame_limit: Option, progress_callback: Option<&dyn Fn(usize, usize)>, -) -> DetectionResults { +) -> anyhow::Result { assert!(opts.lookahead_distance >= 1); - let mut detector = new_detector(dec, opts); - let video_details = y4m::get_video_details(dec); + let mut detector = new_detector::(dec, opts)?; + let video_details = dec.get_video_details()?; let mut frame_queue = BTreeMap::new(); let mut keyframes = BTreeSet::new(); keyframes.insert(0); @@ -229,7 +240,7 @@ pub fn detect_scene_changes( while next_input_frameno < (frameno + opts.lookahead_distance + 1).min(frame_limit.unwrap_or(usize::MAX)) { - let frame = y4m::read_video_frame::(dec, &video_details); + let frame = dec.read_video_frame(&video_details); if let Ok(frame) = frame { frame_queue.insert(next_input_frameno, Arc::new(frame)); next_input_frameno += 1; @@ -275,11 +286,11 @@ pub fn detect_scene_changes( } } } - DetectionResults { + Ok(DetectionResults { scene_changes: keyframes.into_iter().map(|val| val as usize).collect(), frame_count: frameno, speed: frameno as f64 / start_time.elapsed().as_secs_f64(), - } + }) } #[derive(Clone, Copy, Debug, PartialOrd, PartialEq, Eq)] diff --git a/src/main.rs b/src/main.rs index 0a60314..62a8458 100644 --- a/src/main.rs +++ b/src/main.rs @@ -95,7 +95,9 @@ use std::{ }; use anyhow::Result; -use av_scenechange::{detect_scene_changes, DetectionOptions, SceneDetectionSpeed}; +use av_scenechange::{ + decoder::Decoder, detect_scene_changes, DetectionOptions, SceneDetectionSpeed, +}; use clap::Parser; #[derive(Parser, Debug)] @@ -159,12 +161,12 @@ fn main() -> Result<()> { _ => panic!("Speed mode must be in range [0; 1]"), }; - let mut dec = y4m::Decoder::new(&mut reader)?; - let bit_depth = dec.get_bit_depth(); + let mut dec = Decoder::Y4m(y4m::Decoder::new(&mut reader)?); + let bit_depth = dec.get_video_details()?.bit_depth; let results = if bit_depth == 8 { - detect_scene_changes::<_, u8>(&mut dec, opts, None, None) + detect_scene_changes::<_, u8>(&mut dec, opts, None, None)? } else { - detect_scene_changes::<_, u16>(&mut dec, opts, None, None) + detect_scene_changes::<_, u16>(&mut dec, opts, None, None)? }; print!("{}", serde_json::to_string(&results)?); diff --git a/src/vapoursynth.rs b/src/vapoursynth.rs new file mode 100644 index 0000000..6a726d2 --- /dev/null +++ b/src/vapoursynth.rs @@ -0,0 +1,191 @@ +use std::{mem::size_of, path::Path, slice}; + +use anyhow::{bail, ensure}; +use rav1e::{ + color::{ChromaSamplePosition, ChromaSampling}, + data::Rational, + Frame, Pixel, +}; +use vapoursynth::{ + video_info::{Property, VideoInfo}, + vsscript::{Environment, EvalFlags}, +}; + +use crate::y4m::VideoDetails; + +const OUTPUT_INDEX: i32 = 0; + +pub struct VapoursynthDecoder { + env: Environment, + frames_read: usize, + total_frames: usize, +} + +impl VapoursynthDecoder { + /// # Errors + /// + /// - If sourcing an invalid Vapoursynth script. + /// - If using a Vapoursynth script that contains an unsupported video format. + pub fn new(source: &Path) -> anyhow::Result { + let env = Environment::from_file(source, EvalFlags::SetWorkingDir)?; + let total_frames = { + let (node, _) = env.get_output(OUTPUT_INDEX)?; + get_num_frames(node.info())? + }; + Ok(Self { + env, + frames_read: 0, + total_frames, + }) + } + + /// # Errors + /// + /// - If sourcing an invalid Vapoursynth script. + /// - If using a Vapoursynth script that contains an unsupported video format. + pub fn get_video_details(&self) -> anyhow::Result { + let (node, _) = self.env.get_output(OUTPUT_INDEX)?; + let info = node.info(); + let (width, height) = get_resolution(info)?; + Ok(VideoDetails { + width, + height, + bit_depth: get_bit_depth(info)?, + chroma_sampling: get_chroma_sampling(info)?, + chroma_sample_position: ChromaSamplePosition::Unknown, + time_base: get_time_base(info)?, + }) + } + + /// # Errors + /// + /// - If sourcing an invalid Vapoursynth script. + /// - If using a Vapoursynth script that contains an unsupported video format. + /// - If a frame cannot be read. + #[allow(clippy::transmute_ptr_to_ptr)] + pub fn read_video_frame(&mut self, cfg: &VideoDetails) -> anyhow::Result> { + const SB_SIZE_LOG2: usize = 6; + const SB_SIZE: usize = 1 << SB_SIZE_LOG2; + const SUBPEL_FILTER_SIZE: usize = 8; + const FRAME_MARGIN: usize = 16 + SUBPEL_FILTER_SIZE; + const LUMA_PADDING: usize = SB_SIZE + FRAME_MARGIN; + + if self.frames_read >= self.total_frames { + bail!("No frames left"); + } + + let (node, _) = self.env.get_output(OUTPUT_INDEX)?; + let vs_frame = node.get_frame(self.frames_read)?; + self.frames_read += 1; + + let bytes = size_of::(); + let mut f: Frame = + Frame::new_with_padding(cfg.width, cfg.height, cfg.chroma_sampling, LUMA_PADDING); + + // SAFETY: We are using the stride to compute the length of the data slice + unsafe { + f.planes[0].copy_from_raw_u8( + slice::from_raw_parts( + vs_frame.data_ptr(0), + vs_frame.stride(0) * vs_frame.height(0), + ), + vs_frame.stride(0), + bytes, + ); + f.planes[1].copy_from_raw_u8( + slice::from_raw_parts( + vs_frame.data_ptr(1), + vs_frame.stride(1) * vs_frame.height(1), + ), + vs_frame.stride(1), + bytes, + ); + f.planes[2].copy_from_raw_u8( + slice::from_raw_parts( + vs_frame.data_ptr(2), + vs_frame.stride(2) * vs_frame.height(2), + ), + vs_frame.stride(2), + bytes, + ); + } + Ok(f) + } +} + +/// Get the number of frames from a Vapoursynth `VideoInfo` struct. +fn get_num_frames(info: VideoInfo) -> anyhow::Result { + let num_frames = { + if Property::Variable == info.format { + bail!("Cannot output clips with varying format"); + } + if Property::Variable == info.resolution { + bail!("Cannot output clips with varying dimensions"); + } + if Property::Variable == info.framerate { + bail!("Cannot output clips with varying framerate"); + } + + info.num_frames + }; + + ensure!(num_frames != 0, "vapoursynth reported 0 frames"); + + Ok(num_frames) +} + +/// Get the bit depth from a Vapoursynth `VideoInfo` struct. +fn get_bit_depth(info: VideoInfo) -> anyhow::Result { + let bits_per_sample = { + match info.format { + Property::Variable => { + bail!("Cannot output clips with variable format"); + } + Property::Constant(x) => x.bits_per_sample(), + } + }; + + Ok(bits_per_sample as usize) +} + +/// Get the resolution from a Vapoursynth `VideoInfo` struct. +fn get_resolution(info: VideoInfo) -> anyhow::Result<(usize, usize)> { + let resolution = { + match info.resolution { + Property::Variable => { + bail!("Cannot output clips with variable resolution"); + } + Property::Constant(x) => x, + } + }; + + Ok((resolution.width, resolution.height)) +} + +/// Get the time base (inverse of frame rate) from a Vapoursynth `VideoInfo` struct. +fn get_time_base(info: VideoInfo) -> anyhow::Result { + match info.framerate { + Property::Variable => bail!("Cannot output clips with varying framerate"), + Property::Constant(fps) => Ok(Rational::new(fps.denominator, fps.numerator)), + } +} + +/// Get the chroma sampling from a Vapoursynth `VideoInfo` struct. +fn get_chroma_sampling(info: VideoInfo) -> anyhow::Result { + match info.format { + Property::Variable => bail!("Variable pixel format not supported"), + Property::Constant(x) => match x.color_family() { + vapoursynth::format::ColorFamily::YUV => { + let ss = (x.sub_sampling_w(), x.sub_sampling_h()); + match ss { + (1, 1) => Ok(ChromaSampling::Cs420), + (1, 0) => Ok(ChromaSampling::Cs422), + (0, 0) => Ok(ChromaSampling::Cs444), + _ => bail!("Unrecognized chroma subsampling"), + } + } + vapoursynth::format::ColorFamily::Gray => Ok(ChromaSampling::Cs400), + _ => bail!("Currently only YUV input is supported"), + }, + } +}