From 1757ed01938a6a4edb34cdee99f482ca2cb0015f Mon Sep 17 00:00:00 2001
From: SuspiciousDuck <99849016+SuspiciousDuck@users.noreply.github.com>
Date: Sat, 11 Jan 2025 03:39:47 -0600
Subject: [PATCH] Add progress bar to ssimu2 calculation (#7)

---
 auto-boost_2.5.py | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/auto-boost_2.5.py b/auto-boost_2.5.py
index 2849b53..bc1b33b 100755
--- a/auto-boost_2.5.py
+++ b/auto-boost_2.5.py
@@ -1,8 +1,9 @@
 #Originally by Trix
 #Contributors: R1chterScale, Yiss, Kosaka & others from AV1 Weeb edition
 
-from math import ceil
+from math import ceil, floor
 from pathlib import Path
+from tqdm import tqdm
 import json
 import os
 import subprocess
@@ -49,7 +50,7 @@ tmp_dir = Path(args.temp).resolve() if args.temp is not None else output_dir / s
 output_file = output_dir / f"{src_file.stem}_fastpass.mkv"
 scenes_file = tmp_dir / "scenes.json"
 br = float(args.deviation)
-skip = args.skip if args.skip is not None else default_skip
+skip = int(args.skip) if args.skip is not None else default_skip
 aggressive = args.aggressive
 
 def get_ranges(scenes: str) -> list[int]:
@@ -180,7 +181,7 @@ def calculate_ssimu2(src_file, enc_file, ssimu2_txt_path, ranges, skip):
             print(turbo_metrics_run.stdout)
             print(turbo_metrics_run.stderr)
             print("Falling back to vs-zip")
-            skip = args.skip if args.skip is not None else '3'
+            skip = int(args.skip) if args.skip is not None else 3
 
     # If ssimu2zig is True or turbo-metrics failed, use vs-zip
     is_vpy = os.path.splitext(os.path.basename(src_file))[1] == ".vpy"
@@ -199,16 +200,17 @@ def calculate_ssimu2(src_file, enc_file, ssimu2_txt_path, ranges, skip):
     with ssimu2_txt_path.open("w") as file:
         file.write(f"skip: {skip}\n")
     iter = 0
-    for i in range(len(ranges) - 1):
-        cut_source_clip = source_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=1)
-        cut_encoded_clip = encoded_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=1)
-        result = core.vszip.Metrics(cut_source_clip, cut_encoded_clip, mode=0)
-        for index, frame in enumerate(result.frames()):
-            iter += 1
-            score = frame.props['_SSIMULACRA2']
-            with ssimu2_txt_path.open("a") as file:
-                file.write(f"{iter}: {score}\n")
-
+    with tqdm(total=floor(len(source_clip) / int(skip)), desc=f'Calculating SSIMULACRA 2 scores') as pbar:
+        for i in range(len(ranges) - 1):
+            cut_source_clip = source_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=1)
+            cut_encoded_clip = encoded_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=1)
+            result = core.vszip.Metrics(cut_source_clip, cut_encoded_clip, mode=0)
+            for index, frame in enumerate(result.frames()):
+                iter += 1
+                score = frame.props['_SSIMULACRA2']
+                with ssimu2_txt_path.open("a") as file:
+                    file.write(f"{iter}: {score}\n")
+                pbar.update(skip)
 def calculate_xpsnr(src_file, enc_path, xpsnr_txt_path):
     if IS_WINDOWS:
         xpsnr_txt_path = f"{src_file.stem}_xpsnr.log"