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

28
auto-boost_2.5.py
View File

@@ -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"