Fix windows and fallback to vszip if turbo-metrics fails

auto-boost_2.5.py

@@ -9,10 +9,14 @@ import re
 import argparse
 import psutil
 import shutil
+import platform
 import vapoursynth as vs
 core = vs.core
 core.max_cache_size = 1024
 
+IS_WINDOWS = platform.system() == 'Windows'
+NULL_DEVICE = 'NUL' if IS_WINDOWS else '/dev/null'
+
 if shutil.which("av1an") is None:
     raise FileNotFoundError("av1an not found, exiting")
 
@@ -83,6 +87,11 @@ def fast_pass(
     :type workers: int
     """
 
+    # Enclose paths in quotes if they contain spaces
+    input_file = f'"{input_file}"' if ' ' in input_file else input_file
+    output_file = f'"{output_file}"' if ' ' in output_file else output_file
+    tmp_dir = f'"{tmp_dir}"' if ' ' in tmp_dir else tmp_dir
+
     fast_av1an_command = [
         'av1an',
         '-i', input_file,
@@ -147,61 +156,66 @@ def turbo_metrics(
     )
 
 def calculate_ssimu2(src_file, enc_file, ssimu2_txt_path, ranges, skip):
-    if ssimu2zig:
+    if not ssimu2zig:  # Try turbo-metrics first if ssimu2zig is False
-        source_clip = core.lsmas.LWLibavSource(source=src_file, cache=0)
-        encoded_clip = core.lsmas.LWLibavSource(source=enc_file, cache=0)
-
-        #source_clip = source_clip.resize.Bicubic(format=vs.RGBS, matrix_in_s='709').fmtc.transfer(transs="srgb", transd="linear", bits=32)
-        #encoded_clip = encoded_clip.resize.Bicubic(format=vs.RGBS, matrix_in_s='709').fmtc.transfer(transs="srgb", transd="linear", bits=32)
-
-        print(f"source: {len(source_clip)} frames")
-        print(f"encode: {len(encoded_clip)} frames")
-
-        with open(ssimu2_txt_path, "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 open(ssimu2_txt_path, "a") as file:
-                    file.write(f"{iter}: {score}\n")
-    else:
         turbo_metrics_run = turbo_metrics(src_file, enc_file, skip)
-
+        if turbo_metrics_run.returncode == 0:  # If turbo-metrics succeeds
-        if turbo_metrics_run.returncode != 0:
+            with open(ssimu2_txt_path, "w") as file:
+                file.write(f"skip: {skip}\n")
+            frame = 0
+            # for whatever reason, turbo-metrics in csv mode dumps the entire scores to stdout at the end even though it prints them live to stdout.
+            # so we need to see if we've seen ``ssimulacra2`` before and if we have, ignore anything after the second one.
+            ignore_end_barf = False
+            for line in turbo_metrics_run.stdout.splitlines():
+                # set ignore_end_barf to true as this is the first "ssimulacra2" line
+                if line == "ssimulacra2" and not ignore_end_barf:
+                    ignore_end_barf = True
+                # break the loop as we've encountered the second "ssimulacra2" line so we don't get a dupe of the scores.
+                elif line == "ssimulacra2" and ignore_end_barf:
+                    break
+                # assume everything not "ssimulacra2" is a score.
+                if line != "ssimulacra2":
+                    frame += 1
+                    with open(ssimu2_txt_path, "a") as file:
+                        file.write(f"{frame}: {float(line)}\n")
+            return  # Exit if turbo-metrics succeeded
+        else:
             print(f"Turbo Metrics exited with code: {turbo_metrics_run.returncode}")
             print(turbo_metrics_run.stdout)
             print(turbo_metrics_run.stderr)
-            exit(1)
+            print("Falling back to vs-zip")
+            skip = args.skip if args.skip is not None else '3'
 
-        with open(ssimu2_txt_path, "w") as file:
+    # If ssimu2zig is True or turbo-metrics failed, use vs-zip
-            file.write(f"skip: {skip}\n")
+    source_clip = core.lsmas.LWLibavSource(source=src_file, cache=0)
+    encoded_clip = core.lsmas.LWLibavSource(source=enc_file, cache=0)
 
-        frame = 0
+    #source_clip = source_clip.resize.Bicubic(format=vs.RGBS, matrix_in_s='709').fmtc.transfer(transs="srgb", transd="linear", bits=32)
-        # for whatever reason, turbo-metrics in csv mode dumps the entire scores to stdout at the end even though it prints them live to stdout.
+    #encoded_clip = encoded_clip.resize.Bicubic(format=vs.RGBS, matrix_in_s='709').fmtc.transfer(transs="srgb", transd="linear", bits=32)
-        # so we need to see if we've seen ``ssimulacra2`` before and if we have, ignore anything after the second one.
+
-        ignore_end_barf = False
+    print(f"source: {len(source_clip)} frames")
-        for line in turbo_metrics_run.stdout.splitlines():
+    print(f"encode: {len(encoded_clip)} frames")
-            # set ignore_end_barf to true as this is the first "ssimulacra2" line
+    with open(ssimu2_txt_path, "w") as file:
-            if line == "ssimulacra2" and not ignore_end_barf:
+        file.write(f"skip: {skip}\n")
-                ignore_end_barf = True
+    iter = 0
-            # break the loop as we've encountered the second "ssimulacra2" line so we don't get a dupe of the scores.
+    for i in range(len(ranges) - 1):
-            elif line == "ssimulacra2" and ignore_end_barf:
+        cut_source_clip = source_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=1)
-                break
+        cut_encoded_clip = encoded_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=1)
-            # assume everything not "ssimulacra2" is a score.
+        result = core.vszip.Metrics(cut_source_clip, cut_encoded_clip, mode=0)
-            if line != "ssimulacra2":
+        for index, frame in enumerate(result.frames()):
-                frame += 1
+            iter += 1
-                with open(ssimu2_txt_path, "a") as file:
+            score = frame.props['_SSIMULACRA2']
-                    file.write(f"{frame}: {float(line)}\n")
+            with open(ssimu2_txt_path, "a") as file:
+                file.write(f"{iter}: {score}\n")
 
 def calculate_xpsnr(src_file, enc_path, xpsnr_txt_path):
-    xpsnr_command = f'ffmpeg -i {src_file} -i {enc_path} -lavfi xpsnr="stats_file={xpsnr_txt_path}" -f null /dev/null'
+    if IS_WINDOWS:
+        xpsnr_txt_path = xpsnr_txt_path.replace(':', r'\\:')
+
+    if IS_WINDOWS:
+        xpsnr_command = f'ffmpeg -i "{src_file}" -i "{enc_path}" -lavfi xpsnr="stats_file={xpsnr_txt_path}" -f null {NULL_DEVICE}'
+    else:
+        xpsnr_command = f'ffmpeg -i {src_file} -i {enc_path} -lavfi xpsnr="stats_file={xpsnr_txt_path}" -f null {NULL_DEVICE}'
+    
     p = subprocess.Popen(xpsnr_command, shell=True)
     exit_code = p.wait()
 
@@ -481,4 +495,4 @@ match stage:
         calculate_zones(tmp_dir, ranges, zones, crf)
     case _:
         print(f"Stage argument invalid, exiting.")
-        exit(-2)
+        exit(-2)