ffmpeg frame probe to rust

av1an/ffmpeg/ffmpeg.py

@@ -7,37 +7,6 @@ from subprocess import PIPE, STDOUT
 from typing import List
 from av1an.logger import log
 
-# TODO: redo to module, add ffmpeg scenedetection for fallback
-
-
-def frame_probe_ffmpeg(source: Path):
-    """
-    Get frame count.
-    Direct counting of frame count using ffmpeg. Slow, Precise.
-    :param: source: Path to input file
-    """
-    cmd = [
-        "ffmpeg",
-        "-hide_banner",
-        "-i",
-        source.as_posix(),
-        "-map",
-        "0:v:0",
-        "-c",
-        "copy",
-        "-f",
-        "null",
-        "-",
-    ]
-    r = subprocess.run(cmd, stdout=PIPE, stderr=PIPE)
-    matches = re.findall(
-        r"frame=\s*([0-9]+)\s", r.stderr.decode("utf-8") + r.stdout.decode("utf-8")
-    )
-    total = int(matches[-1])
-    log("Get frame count with ffmpeg")
-    log(f"Frame count: {total}")
-    return total
-
 
 def get_frametypes(file: Path) -> List:
     """

av1an/utils.py

@@ -1,16 +1,14 @@
 #!/bin/env python
 
 import re
-import sys
 from typing import List
 from pathlib import Path
 import cv2
 import numpy as np
-import hashlib
 
-from av1an.ffmpeg import frame_probe_ffmpeg
 from av1an.vapoursynth import frame_probe_vspipe, is_vapoursynth
 from av1an.logger import log
+from av1an.av1an import ffmpeg_get_frame_count
 
 
 def list_index_of_regex(lst: List[str], regex_str: str) -> int:
@@ -71,4 +69,4 @@ def frame_probe(source: Path):
     if is_vapoursynth(source):
         return frame_probe_vspipe(source)
 
-    return frame_probe_ffmpeg(source)
+    return ffmpeg_get_frame_count(str(source.resolve()))

src/lib.rs

@@ -1,12 +1,56 @@
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
-use std::process::{Command, Stdio};
+use regex::Regex;
 use std::{
   collections::hash_map::DefaultHasher,
   hash::{Hash, Hasher},
+  path::Path,
+};
+use std::{
+  io::{self, stderr, Write},
+  os::unix::prelude::CommandExt,
+};
+use std::{
+  process::{Command, Stdio},
+  str,
 };
 
-/// Formats the sum of two numbers as string.
+#[pyfunction]
+/// Get frame count. Direct counting of frame count using ffmpeg
+fn ffmpeg_get_frame_count(source: String) -> PyResult<usize> {
+  let source_path = Path::new(&source);
+
+  let mut cmd = Command::new("ffmpeg");
+  cmd.args(&[
+    "-hide_banner",
+    "-i",
+    source_path.to_str().unwrap(),
+    "-map",
+    "0:v:0",
+    "-c",
+    "copy",
+    "-f",
+    "null",
+    "-",
+  ]);
+
+  cmd.stdout(Stdio::piped());
+  cmd.stderr(Stdio::piped());
+
+  let out = cmd.output()?;
+  assert!(out.status.success());
+
+  let re = Regex::new(r".*frame=\s*([0-9]+)\s").unwrap();
+  let output = String::from_utf8(out.stderr).unwrap();
+
+  // io::stdout().write_all(output.as_bytes()).unwrap();
+  // dbg!(&output);
+  let cap = re.captures(&output).unwrap();
+
+  let frame_count = cap[cap.len() - 1].parse::<usize>().unwrap();
+  Ok(frame_count)
+}
+
 #[pyfunction]
 fn get_ffmpeg_info() -> PyResult<String> {
   let mut cmd = Command::new("ffmpeg");
@@ -42,6 +86,7 @@ fn av1an(_py: Python, m: &PyModule) -> PyResult<()> {
   m.add_function(wrap_pyfunction!(get_ffmpeg_info, m)?)?;
   m.add_function(wrap_pyfunction!(hash_path, m)?)?;
   m.add_function(wrap_pyfunction!(adapt_probing_rate, m)?)?;
+  m.add_function(wrap_pyfunction!(ffmpeg_get_frame_count, m)?)?;
 
   Ok(())
 }