Port tqdm to Rust with indicatif crate (#284)

Port progress bar to Rust with indicatif

README.md

@@ -61,7 +61,7 @@ With your own parameters:
 
     --keep                  Doesn't delete temporary folders after encode has finished.
 
-    -q --quiet              Do not print tqdm to the terminal.
+    -q --quiet              Do not print a progress bar to the terminal.
 
     -l --logging          Path to .log file(By default created in temp folder)
 

appveyor.yml

@@ -3,7 +3,7 @@ environment:
   PYTHON: C:\Python38-x64
 
 install:
-- cmd: '"%PYTHON%\python.exe" -m pip install numpy opencv-python tqdm scenedetect[opencv,progress_bar] psutil pypiwin32 '
+- cmd: '"%PYTHON%\python.exe" -m pip install numpy opencv-python scenedetect[opencv,progress_bar] psutil pypiwin32 '
 - cmd: '"%PYTHON%\python.exe" -m pip install pyinstaller scipy matplotlib'
 
 build_script:

av1an-core/src/lib.rs

@@ -207,7 +207,7 @@ struct Baz {
 }
 
 pub fn read_file_to_string(file: PathBuf) -> Result<String, Error> {
-  Ok(fs::read_to_string(&file).unwrap_or_else(|_| panic!(" Can't open file {:?}", file)))
+  Ok(fs::read_to_string(&file).unwrap_or_else(|_| panic!("Can't open file {:?}", file)))
 }
 
 pub fn read_weighted_vmaf(file: PathBuf, percentile: f64) -> Result<f64, serde_json::Error> {

av1an-encoder-constructor/src/lib.rs

@@ -4,6 +4,7 @@ use std::borrow::Cow;
 use std::cmp;
 use std::path::PathBuf;
 use std::{str::FromStr, usize};
+
 macro_rules! into_vec {
   ($($x:expr),* $(,)?) => {
     vec![
@@ -216,31 +217,24 @@ impl Encoder {
 
   pub fn compose_2_2_pass(&self, params: Vec<String>, fpf: String, output: String) -> Vec<String> {
     match &self {
-      // Aomenc
       Self::aom => chain!(
         into_vec!["aomenc", "--passes=2", "--pass=2"],
         params,
         into_vec![format!("--fpf={}.log", fpf), "-o", output, "-"],
       )
       .collect(),
-
-      // Rav1e
       Self::rav1e => chain!(
         into_vec!["rav1e", "-", "-y", "-q"],
         params,
         into_vec!["--second-pass", format!("{}.stat", fpf), "--output", output,]
       )
       .collect(),
-
-      // VPX
       Self::libvpx => chain!(
         into_vec!["vpxenc", "--passes=2", "--pass=2"],
         params,
         into_vec![format!("--fpf={}.log", fpf), "-o", output, "-"],
       )
       .collect(),
-
-      // SVT-AV1
       Self::svt_av1 => chain!(
         into_vec![
           "SvtAv1EncApp",
@@ -262,8 +256,6 @@ impl Encoder {
         ],
       )
       .collect(),
-
-      // x264
       Self::x264 => chain!(
         into_vec![
           "x264",
@@ -279,8 +271,6 @@ impl Encoder {
         into_vec!["--stats", format!("{}.log", fpf), "-", "-o", output,]
       )
       .collect(),
-
-      // x265
       Self::x265 => chain!(
         into_vec![
           "x265",
@@ -301,7 +291,6 @@ impl Encoder {
 
   pub fn get_default_arguments(&self) -> Vec<&str> {
     match &self {
-      // Aomenc
       Encoder::aom => into_vec![
         "--threads=8",
         "-b",
@@ -312,8 +301,6 @@ impl Encoder {
         "--tile-columns=2",
         "--tile-rows=1",
       ],
-
-      // Rav1e
       Encoder::rav1e => into_vec![
         "--tiles",
         "8",
@@ -323,8 +310,6 @@ impl Encoder {
         "100",
         "--no-scene-detection",
       ],
-
-      //VPX
       Encoder::libvpx => into_vec![
         "--codec=vp9",
         "-b",
@@ -336,16 +321,8 @@ impl Encoder {
         "--cq-level=30",
         "--row-mt=1",
       ],
-
-      // SVT-AV1
-      Encoder::svt_av1 => {
-        into_vec!["--preset", "4", "--keyint", "240", "--rc", "0", "--crf", "25",]
-      }
-
-      // x264
+      Encoder::svt_av1 => into_vec!["--preset", "4", "--keyint", "240", "--rc", "0", "--crf", "25"],
       Encoder::x264 => into_vec!["--preset", "slow", "--crf", "25"],
-
-      // x265
       Encoder::x265 => into_vec!["-p", "slow", "--crf", "25", "-D", "10"],
     }
   }
@@ -450,7 +427,7 @@ impl Encoder {
     }
   }
 
-  pub fn match_line(&self, line: String) -> Option<usize> {
+  pub fn match_line(&self, line: &str) -> Option<usize> {
     let encoder_regex = Regex::new(self.pipe_match()).unwrap();
     if !encoder_regex.is_match(&line) {
       return Some(0);
@@ -628,6 +605,7 @@ impl Encoder {
       ],
     }
   }
+
   pub fn probe_cmd(
     &self,
     temp: String,
@@ -675,6 +653,7 @@ impl Encoder {
 
     (pipe, output)
   }
+
   pub fn construct_target_quality_slow_command(&self, q: String) -> Vec<Cow<str>> {
     match &self {
       Encoder::aom => into_vec!["aomenc", "--passes=1", format!("--cq-level={}", q),],

av1an-pyo3/Cargo.toml

@@ -13,7 +13,6 @@ requires-dist = [
     "numpy",
     "scenedetect[opencv]",
     "opencv-python",
-    "tqdm",
     "psutil",
     "scipy",
     "matplotlib",
@@ -32,6 +31,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 pyo3 = { version="0.13.2" }
+indicatif = { git = "https://github.com/redzic/indicatif", branch = "fps-template" }
 once_cell = "1.8.0"
 chrono = "0.4.19"
 av1an-core = { path="../av1an-core" }

av1an-pyo3/pyproject.toml

@@ -15,7 +15,6 @@ requires-dist = [
     "numpy",
     "scenedetect[opencv]",
     "opencv-python",
-    "tqdm",
     "psutil",
     "scipy",
     "matplotlib",

av1an-pyo3/src/lib.rs

@@ -1,3 +1,5 @@
+use indicatif::ProgressBar;
+use indicatif::ProgressStyle;
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
 
@@ -361,7 +363,7 @@ fn man_command(encoder: String, params: Vec<String>, q: usize) -> Vec<String> {
 }
 
 #[pyfunction]
-fn match_line(encoder: String, line: String) -> PyResult<usize> {
+fn match_line(encoder: &str, line: &str) -> PyResult<usize> {
   let enc = av1an_encoder_constructor::Encoder::from_str(&encoder).unwrap();
 
   Ok(enc.match_line(line).unwrap())
@@ -400,8 +402,45 @@ pub fn read_weighted_vmaf(fl: String, percentile: f64) -> PyResult<f64> {
   Ok(val)
 }
 
+const INDICATIF_PROGRESS_TEMPLATE: &str = "{spinner:.green} [{elapsed_precise}] [{bar:60.cyan/blue}] {percent:>3.bold}% {pos}/{len} ({fps:.bold}, eta {eta})";
+
+static PROGRESS_BAR: OnceCell<ProgressBar> = OnceCell::new();
+
+#[pyfunction]
+pub fn init_progress_bar(len: u64) {
+  PROGRESS_BAR.get_or_init(|| {
+    let bar = ProgressBar::new(len);
+    bar.set_style(
+      ProgressStyle::default_bar()
+        .template(INDICATIF_PROGRESS_TEMPLATE)
+        .progress_chars("#>-"),
+    );
+    bar
+  });
+}
+
+#[pyfunction]
+pub fn update_bar(inc: u64) {
+  PROGRESS_BAR
+    .get()
+    .expect("The progress bar was not initialized!")
+    .inc(inc)
+}
+
+#[pyfunction]
+pub fn finish_progress_bar() {
+  PROGRESS_BAR
+    .get()
+    .expect("The progress bar was not initialized!")
+    .finish();
+}
+
 #[pymodule]
 fn av1an_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
+  m.add_function(wrap_pyfunction!(init_progress_bar, m)?)?;
+  m.add_function(wrap_pyfunction!(update_bar, m)?)?;
+  m.add_function(wrap_pyfunction!(finish_progress_bar, m)?)?;
+
   m.add_function(wrap_pyfunction!(get_ffmpeg_info, m)?)?;
   m.add_function(wrap_pyfunction!(determine_workers, m)?)?;
   m.add_function(wrap_pyfunction!(create_vs_file, m)?)?;

av1an-scene-detection/src/aom_kf.rs

@@ -39,7 +39,7 @@ pub struct AomFirstPassStats {
 }
 
 pub fn read_aomenc_stats_struct(file: PathBuf) -> Vec<AomFirstPassStats> {
-  let raw_data: Vec<u8> = read(file).unwrap();
+  let mut raw_data: Vec<u8> = read(file).unwrap();
   let frame_list: Vec<AomFirstPassStats> = unsafe { transmute(raw_data) };
   frame_list
 }
@@ -70,9 +70,9 @@ fn test_candidate_kf(
   current_frame_index: u64,
   frame_count_so_far: u64,
 ) -> bool {
-  let p: AomFirstPassStats = dict_list[(current_frame_index - 1) as usize];
-  let c: AomFirstPassStats = dict_list[(current_frame_index) as usize];
-  let f: AomFirstPassStats = dict_list[(current_frame_index + 1) as usize];
+  let p = dict_list[(current_frame_index - 1) as usize];
+  let c = dict_list[(current_frame_index) as usize];
+  let f = dict_list[(current_frame_index + 1) as usize];
 
   let boost_factor = 12.5;
   let min_intra_level = 0.25;
@@ -108,7 +108,7 @@ fn test_candidate_kf(
     let mut old_boost_score = 0.0;
     let mut decay_accumulator = 1.0;
     for i in 0..16 {
-      let lnf: AomFirstPassStats = dict_list[(current_frame_index + 1 + i) as usize];
+      let lnf = dict_list[(current_frame_index + 1 + i) as usize];
       let mut next_iiratio = boost_factor * lnf.intra_error / double_divide_check(lnf.coded_error);
 
       if next_iiratio > kf_ii_max {
@@ -136,7 +136,7 @@ fn test_candidate_kf(
       }
       old_boost_score = boost_score;
 
-      // If there is tolerable prediction for at least the next 3 frames then break out else discard this potential key frame && move on
+      // If there is tolerable prediction for at least the next 3 frames then break out else discard this potential key frame and move on
       if boost_score > 30.0 && (i > 3) {
         is_keyframe = true;
       }

av1an/manager/Counter.py

@@ -1,9 +1,5 @@
 from multiprocessing.managers import BaseManager
-
-try:
-    from tqdm import tqdm
-except ImportError:
-    tqdm = None
+from av1an_pyo3 import update_bar, init_progress_bar, finish_progress_bar
 
 
 def Manager():
@@ -16,34 +12,22 @@ def Manager():
 
 
 class Counter:
-    def __init__(self, total, initial, use_tqdm=True):
+    def __init__(self, total, initial, quiet=False):
         self.first_update = True
         self.initial = initial
         self.left = total - initial
         self.current = 0
-        self.use_tqdm = use_tqdm and (tqdm is not None)
-        if use_tqdm:
-            self.tqdm_bar = tqdm(
-                total=self.left,
-                initial=0,
-                dynamic_ncols=True,
-                unit="fr",
-                leave=True,
-                smoothing=0.01,
-            )
+        self.quiet = quiet
+        if not self.quiet:
+            init_progress_bar(total)
 
     def update(self, value):
-        if self.use_tqdm:
-            if self.first_update:
-                self.tqdm_bar.reset(self.left)
-                self.first_update = False
-            self.tqdm_bar.update(value)
-        else:
-            self.current += value
+        if not self.quiet:
+            update_bar(value)
 
     def close(self):
-        if self.use_tqdm:
-            self.tqdm_bar.close()
+        if not self.quiet:
+            finish_progress_bar()
 
     def get_frames(self):
         return self.current

av1an/manager/Manager.py

@@ -144,6 +144,6 @@ def startup(self, project: Project, chunk_queue: List[Chunk]):
         )
         BaseManager.register("Counter", Counter)
         counter = Manager().Counter(
-            project.get_frames(), self.initial_frames, not project.quiet
+            project.get_frames(), self.initial_frames, project.quiet
         )
         project.counter = counter

av1an/manager/Pipes.py

@@ -85,7 +85,7 @@ def process_encoding_pipe(
         raise RuntimeError("Error in processing encoding pipe").with_traceback(tb)
 
 
-def tqdm_bar(
+def create_pipes(
     a: Project, c: Chunk, encoder, counter, frame_probe_source, passes, current_pass
 ):
 

av1an/manager/Queue.py

@@ -4,14 +4,13 @@
 import sys
 import time
 import traceback
-from pathlib import Path
 
 from av1an.chunk import Chunk
 from av1an.target_quality import TargetQuality
 from av1an.utils import frame_probe
 from av1an_pyo3 import log
 
-from .Pipes import tqdm_bar
+from .Pipes import create_pipes
 
 
 class Queue:
@@ -59,7 +58,7 @@ def encode_chunk(self, chunk: Chunk):
 
                 # Run all passes for this chunk
                 for current_pass in range(1, self.project.passes + 1):
-                    tqdm_bar(
+                    create_pipes(
                         self.project,
                         chunk,
                         self.project.encoder,

requirements.txt

@@ -1,7 +1,6 @@
 numpy
 scenedetect[opencv]
 opencv-python
-tqdm
 psutil
 scipy
 matplotlib

setup.py

@@ -20,7 +20,6 @@
     "vapoursynth",
     "scenedetect[opencv]",
     "opencv-python",
-    "tqdm",
     "psutil",
     "scipy",
     "matplotlib",