extra_splits to rust

av1an-core/src/lib.rs

@@ -1,4 +1,5 @@
 #![allow(unused)]
+#![feature(iter_zip)]
 
 #[macro_use]
 extern crate log;
@@ -15,6 +16,7 @@ use sysinfo::SystemExt;
 
 pub mod concat;
 pub mod ffmpeg;
+pub mod split;
 pub mod target_quality;
 pub mod vapoursynth;
 

av1an-core/src/split.rs

@@ -0,0 +1,32 @@
+use std::iter::zip;
+
+pub fn extra_splits(
+  split_locations: Vec<usize>,
+  total_frames: usize,
+  split_size: usize,
+) -> Vec<usize> {
+  let mut result_vec: Vec<usize> = split_locations.clone();
+
+  let mut total_length = split_locations.clone();
+  total_length.insert(0, 0);
+  total_length.push(total_frames);
+
+  let iter = total_length[..total_length.len() - 1]
+    .iter()
+    .zip(total_length[1..].iter());
+
+  for (x, y) in iter {
+    let distance = y - x;
+    if distance > split_size {
+      let additional_splits = (distance / split_size) + 1;
+      for n in 1..additional_splits {
+        let new_split = (distance as f64 * (n as f64 / additional_splits as f64)) as usize + x;
+        result_vec.push(new_split);
+      }
+    }
+  }
+
+  result_vec.sort();
+
+  result_vec
+}

av1an-pyo3/src/lib.rs

@@ -166,6 +166,11 @@ fn concatenate_ffmpeg(temp: String, output: String, encoder: String) -> PyResult
   ))
 }
 
+#[pyfunction]
+fn extra_splits(split_locations: Vec<usize>, total_frames: usize, split_size: usize) -> Vec<usize> {
+  av1an_core::split::extra_splits(split_locations, total_frames, split_size)
+}
+
 #[pymodule]
 fn av1an_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
   m.add_function(wrap_pyfunction!(get_ffmpeg_info, m)?)?;
@@ -181,5 +186,6 @@ fn av1an_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
   m.add_function(wrap_pyfunction!(concatenate_ffmpeg, m)?)?;
   m.add_function(wrap_pyfunction!(extract_audio, m)?)?;
   m.add_function(wrap_pyfunction!(get_frame_types, m)?)?;
+  m.add_function(wrap_pyfunction!(extra_splits, m)?)?;
   Ok(())
 }

av1an/split.py

@@ -13,6 +13,8 @@ from .project import Project
 from .scenedetection import aom_keyframes, AOM_KEYFRAMES_DEFAULT_PARAMS, pyscene, ffmpeg
 from .logger import log
 
+from av1an_pyo3 import extra_splits
+
 # TODO: organize to single segmenting/splitting module
 
 
@@ -54,7 +56,10 @@ def split_routine(project: Project, resuming: bool) -> List[int]:
 
     # Applying extra splits
     if project.extra_split:
-        scenes = extra_splits(project, scenes)
+        log("Applying extra splits")
+        log(f"Split distance: {project.extra_split}")
+        scenes = extra_splits(scenes, project.get_frames(), project.extra_split)
+        log(f"New splits:{len(scenes)}")
 
     # write scenes for resuming later if needed
     return scenes
@@ -130,31 +135,6 @@ def segment(video: Path, temp: Path, frames: List[int]):
     log("Split Done")
 
 
-def extra_splits(project: Project, split_locations: list):
-    log("Applying extra splits")
-
-    split_locs_with_start = split_locations[:]
-    split_locs_with_start.insert(0, 0)
-
-    split_locs_with_end = split_locations[:]
-    split_locs_with_end.append(project.get_frames())
-
-    splits = list(zip(split_locs_with_start, split_locs_with_end))
-    for i in splits:
-        distance = i[1] - i[0]
-        if distance > project.extra_split:
-            to_add = distance // project.extra_split
-            new_scenes = list(
-                linspace(i[0], i[1], to_add + 1, dtype=int, endpoint=False)[1:]
-            )
-            split_locations.extend(new_scenes)
-
-    result = [int(x) for x in sorted(split_locations)]
-    log(f"Split distance: {project.extra_split}")
-    log(f"New splits:{len(result)}")
-    return result
-
-
 def calc_split_locations(project: Project) -> List[int]:
     """
     Determines a list of frame numbers to split on with pyscenedetect or aom keyframes