pat-e/auto-boost-algorithm
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add progress bar to ssimu2 calculation (#7)

Browse Source
This commit is contained in:
SuspiciousDuck
2025-01-11 03:39:47 -06:00
committed by GitHub
parent 2e55a3eaae
commit 1757ed0193
1 changed files with 15 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
auto-boost_2.5.py
View File

@@ -1,8 +1,9 @@
#Originally by Trix #Originally by Trix
#Contributors: R1chterScale, Yiss, Kosaka & others from AV1 Weeb edition #Contributors: R1chterScale, Yiss, Kosaka & others from AV1 Weeb edition
from math import ceil from math import ceil, floor
from pathlib import Path from pathlib import Path
from tqdm import tqdm
import json import json
import os import os
import subprocess 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" output_file = output_dir / f"{src_file.stem}_fastpass.mkv"
scenes_file = tmp_dir / "scenes.json" scenes_file = tmp_dir / "scenes.json"
br = float(args.deviation) 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 aggressive = args.aggressive
def get_ranges(scenes: str) -> list[int]: 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.stdout)
print(turbo_metrics_run.stderr) print(turbo_metrics_run.stderr)
print("Falling back to vs-zip") 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 # If ssimu2zig is True or turbo-metrics failed, use vs-zip
is_vpy = os.path.splitext(os.path.basename(src_file))[1] == ".vpy" is_vpy = os.path.splitext(os.path.basename(src_file))[1] == ".vpy"
@@ -199,6 +200,7 @@ def calculate_ssimu2(src_file, enc_file, ssimu2_txt_path, ranges, skip):
with ssimu2_txt_path.open("w") as file: with ssimu2_txt_path.open("w") as file:
file.write(f"skip: {skip}\n") file.write(f"skip: {skip}\n")
iter = 0 iter = 0
with tqdm(total=floor(len(source_clip) / int(skip)), desc=f'Calculating SSIMULACRA 2 scores') as pbar:
for i in range(len(ranges) - 1): for i in range(len(ranges) - 1):
cut_source_clip = source_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=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) cut_encoded_clip = encoded_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=1)
@@ -208,7 +210,7 @@ def calculate_ssimu2(src_file, enc_file, ssimu2_txt_path, ranges, skip):
score = frame.props['_SSIMULACRA2'] score = frame.props['_SSIMULACRA2']
with ssimu2_txt_path.open("a") as file: with ssimu2_txt_path.open("a") as file:
file.write(f"{iter}: {score}\n") file.write(f"{iter}: {score}\n")
pbar.update(skip)
def calculate_xpsnr(src_file, enc_path, xpsnr_txt_path): def calculate_xpsnr(src_file, enc_path, xpsnr_txt_path):
if IS_WINDOWS: if IS_WINDOWS:
xpsnr_txt_path = f"{src_file.stem}_xpsnr.log" xpsnr_txt_path = f"{src_file.stem}_xpsnr.log"