diff --git a/auto-boost_2.5.py b/auto-boost_2.5.py new file mode 100755 index 0000000..a4cc499 --- /dev/null +++ b/auto-boost_2.5.py @@ -0,0 +1,484 @@ +#Originally by Trix +#Contributors: R1chterScale, Yiss and Kosaka + +from math import ceil +import json +import os +import subprocess +import re +import argparse +import psutil +import shutil +import vapoursynth as vs +core = vs.core +core.max_cache_size = 1024 + +if shutil.which("av1an") is None: + raise FileNotFoundError("av1an not found, exiting") + +if shutil.which("turbo-metrics") is None: + print("turbo-metrics not found, defaulting to vs-zip") + ssimu2zig = True + default_skip = 3 +else: + ssimu2zig = False + default_skip = 0 + +parser = argparse.ArgumentParser() +parser.add_argument("-s", "--stage", help = "Select stage: 1 = encode, 2 = calculate metrics, 3 = generate zones | Default: all", default=0) +parser.add_argument("-i", "--input", required=True, help = "Video input filepath (original source file)") +parser.add_argument("-q", "--quality", help = "Base quality (CRF) | Default: 30", default=30) +parser.add_argument("-d", "--deviation", help = "Maximum CRF change from original | Default: 10", default=10) +parser.add_argument("-p", "--preset", help = "Fast encode preset | Default: 9", default=9) +parser.add_argument("-w", "--workers", help = "Number of av1an workers | Default: amount of physical cores", default=psutil.cpu_count(logical=False)) +parser.add_argument("-m", "--metrics", help = "Select metrics: 1 = SSIMU2, 2 = XPSNR, 3 = Both | Default: 1", default=1) +parser.add_argument("-S", "--skip", help = "SSIMU2 skip value, every nth frame's SSIMU2 is calculated | Default: 0 for turbo-metrics, 3 for vs-zip") +parser.add_argument("-z", "--zones", help = "Zones calculation method: 1 = SSIMU2, 2 = XPSNR, 3 = Multiplication, 4 = Lowest Result | Default: 1", default=1) +parser.add_argument("-a", "--aggressive", action='store_true', help = "More aggressive boosting | Default: not active") +args = parser.parse_args() +stage = int(args.stage) +src_file = args.input +output_dir = os.path.dirname(src_file) +tmp_dir = os.path.join(output_dir, "temp") +output_file = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_fastpass.mkv") +scenes_file = os.path.join(tmp_dir, "scenes.json") +br = float(args.deviation) +skip = args.skip if args.skip is not None else default_skip +aggressive = args.aggressive + +def get_ranges(scenes: str) -> list[int]: + """ + Reads a scene file and returns a list of frame numbers for each scene change. + + :param scenes: path to scene file + :type scenes: str + + :return: list of frame numbers + :rtype: list[int] + """ + ranges = [0] + with open(scenes, "r") as file: + content = json.load(file) + for scene in content['scenes']: + ranges.append(scene['end_frame']) + return ranges + +def fast_pass( + input_file: str, output_file: str, tmp_dir: str, preset: int, crf: float, workers: int +): + """ + Quick fast-pass using Av1an + + :param input_file: path to input file + :type input_file: str + :param output_file: path to output file + :type output_file: str + :param tmp_dir: path to temporary directory + :type tmp_dir: str + :param preset: encoder preset + :type preset: int + :param crf: target CRF + :type crf: float + :param workers: number of workers + :type workers: int + """ + + fast_av1an_command = [ + 'av1an', + '-i', input_file, + '--temp', tmp_dir, + '-y', + '--verbose', + '--keep', + '-m', 'lsmash', + '-c', 'mkvmerge', + '--min-scene-len', '24', + '--sc-downscale-height', '720', + '--set-thread-affinity', '2', + '-e', 'svt-av1', + '--force', + '-v', f'"--preset {preset} --crf {crf:.2f} --lp 2 --scm 0 --keyint 0 --fast-decode 1 --color-primaries 1 --transfer-characteristics 1 --matrix-coefficients 1"', + '-w', str(workers), + '-o', output_file + ] + + process = subprocess.run(' '.join(fast_av1an_command), shell=True, check=True) + + if process.returncode != 0: + print(f"Av1an exited with code: {process.returncode}") + exit(1) + +def turbo_metrics( + source: str, distorted: str, every: int +) -> subprocess.CompletedProcess: + """ + Compare two files with SSIMULACRA2 using turbo-metrics. + + :param source: path to source file + :type source: str + :param distorted: path to distorted file + :type distorted: str + :param every: compare every X frames + :type every: int + + :return: completed process + :rtype: subprocess.CompletedProcess + """ + + turbo_cmd = [ + "turbo-metrics", + "-m", + "ssimulacra2", + "--output", + "csv", + ] + + if every > 1: + turbo_cmd.append("--every") + turbo_cmd.append(str(every)) + + turbo_cmd.append(source) + turbo_cmd.append(distorted) + + return subprocess.run( + turbo_cmd, + capture_output=True, + text=True, + ) + +def calculate_ssimu2(src_file, enc_file, ssimu2_txt_path, ranges, skip): + if ssimu2zig: + 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: + print(f"Turbo Metrics exited with code: {turbo_metrics_run.returncode}") + print(turbo_metrics_run.stdout) + print(turbo_metrics_run.stderr) + exit(1) + + 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") + +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' + p = subprocess.Popen(xpsnr_command, shell=True) + exit_code = p.wait() + + if exit_code != 0: + print("XPSNR encountered an error, exiting.") + exit(-2) + +def get_xpsnr(xpsnr_txt_path): + count=0 + + sum_weighted = 0 + values_weighted: list[int] = [] + + with open(xpsnr_txt_path, "r") as file: + for line in file: + match = re.search(r"XPSNR Y: ([0-9]+\.[0-9]+) XPSNR U: ([0-9]+\.[0-9]+) XPSNR : ([0-9]+\.[0-9]+)", line) + if match: + Y = float(match.group(1)) + U = float(match.group(2)) + V = float(match.group(3)) + W = (4 * Y + U + V) / 6 + + sum_weighted += W + values_weighted.append(W) + + count += 1 + else: + print(line) + avg_weighted = sum_weighted / count + for i in range(len(values_weighted)): + values_weighted[i] /= avg_weighted + return values_weighted + +def get_ssimu2(ssimu2_txt_path): + ssimu2_scores: list[int] = [] + + with open(ssimu2_txt_path, "r") as file: + skipmatch = re.search(r"skip: ([0-9]+)", file.readline()) + if skipmatch: + skip = int(skipmatch.group(1)) + else: + print("Skip value not detected in SSIMU2 file, exiting.") + exit(-2) + for line in file: + match = re.search(r"([0-9]+): ([0-9]+\.[0-9]+)", line) + if match: + score = float(match.group(2)) + ssimu2_scores.append(score) + else: + print(line) + return ssimu2_scores, skip + +def calculate_std_dev(score_list: list[int]): + """ + Takes a list of metrics scores and returns the associated arithmetic mean, + 5th percentile and 95th percentile scores. + + :param score_list: list of SSIMU2 scores + :type score_list: list + """ + + filtered_score_list = [score for score in score_list if score >= 0] + sorted_score_list = sorted(filtered_score_list) + average = sum(filtered_score_list)/len(filtered_score_list) + percentile_5 = sorted_score_list[len(filtered_score_list)//20] + percentile_95 = sorted_score_list[int (len(filtered_score_list)//(20/19))] + return (average, percentile_5, percentile_95) + +def generate_zones(ranges: list, percentile_5_total: list, average: int, crf: float, zones_txt_path: str): + """ + Appends a scene change to the ``zones_txt_path`` file in Av1an zones format. + + creates ``zones_txt_path`` if it does not exist. If it does exist, the line is + appended to the end of the file. + + :param ranges: Scene changes list + :type ranges: list + :param percentile_5_total: List containing all 5th percentile scores + :type percentile_5_total: list + :param average: Full clip average score + :type average: int + :param crf: CRF setting to use for the zone + :type crf: int + :param zones_txt_path: Path to the zones.txt file + :type zones_txt_path: str + """ + zones_iter = 0 + for i in range(len(ranges)-1): + zones_iter += 1 + if aggressive: + new_crf = crf - ceil((1.0 - (percentile_5_total[i] / average)) * 40 * 4) / 4 + else: + new_crf = crf - ceil((1.0 - (percentile_5_total[i] / average)) * 20 * 4) / 4 + + if new_crf < crf - br: # set lowest allowed crf + new_crf = crf - br + + if new_crf > crf + br: # set highest allowed crf + new_crf = crf + br + + print(f'Enc: [{ranges[i]}:{ranges[i+1]}]\n' + f'Chunk 5th percentile: {percentile_5_total[i]}\n' + f'Adjusted CRF: {new_crf:.2f}\n') + + with open(zones_txt_path, "w" if zones_iter == 1 else "a") as file: + file.write(f"{ranges[i]} {ranges[i+1]} svt-av1 --crf {new_crf:.2f}\n") + +def calculate_metrics(src_file, output_file, tmp_dir, ranges, skip, metrics): + match metrics: + case 1: + ssimu2_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_ssimu2.log") + calculate_ssimu2(src_file, output_file, ssimu2_txt_path, ranges, skip) + case 2: + xpsnr_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_xpsnr.log") + calculate_xpsnr(src_file, output_file, xpsnr_txt_path) + case 3: + xpsnr_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_xpsnr.log") + ssimu2_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_ssimu2.log") + calculate_xpsnr(src_file, output_file, xpsnr_txt_path) + calculate_ssimu2(src_file, output_file, ssimu2_txt_path, ranges, skip) + +def calculate_zones(tmp_dir, ranges, zones, cq): + match zones: + case 1: + ssimu2_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_ssimu2.log") + (ssimu2_scores, skip) = get_ssimu2(ssimu2_txt_path) + ssimu2_zones_txt_path = f"{tmp_dir}/ssimu2_zones.txt" + ssimu2_total_scores: list[int] = [] + ssimu2_percentile_5_total = [] + ssimu2_iter = 0 + + for i in range(len(ranges)-1): + ssimu2_chunk_scores: list[int] = [] + xpsnr_chunk_scores: list[int] = [] + ssimu2_frames = (ranges[i+1] - ranges[i]) // skip + for frames in range(ssimu2_frames): + ssimu2_score = ssimu2_scores[ssimu2_iter] + ssimu2_chunk_scores.append(ssimu2_score) + ssimu2_total_scores.append(ssimu2_score) + ssimu2_iter += 1 + (ssimu2_average, ssimu2_percentile_5, ssimu2_percentile_95) = calculate_std_dev(ssimu2_chunk_scores) + ssimu2_percentile_5_total.append(ssimu2_percentile_5) + #print(f'5th Percentile: {ssimu2_percentile_5}') + (ssimu2_average, ssimu2_percentile_5, ssimu2_percentile_95) = calculate_std_dev(ssimu2_total_scores) + + print(f'SSIMU2:') + print(f'Median score: {ssimu2_average}') + print(f'5th Percentile: {ssimu2_percentile_5}') + print(f'95th Percentile: {ssimu2_percentile_95}\n') + generate_zones(ranges, ssimu2_percentile_5_total, ssimu2_average, cq, ssimu2_zones_txt_path) + + case 2: + xpsnr_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_xpsnr.log") + xpsnr_scores: list[int] = get_xpsnr(xpsnr_txt_path) + xpsnr_zones_txt_path = f"{tmp_dir}/xpsnr_zones.txt" + xpsnr_total_scores: list[int] = [] + xpsnr_percentile_5_total = [] + xpsnr_iter = 0 + + for i in range(len(ranges)-1): + xpsnr_chunk_scores: list[int] = [] + xpsnr_frames = (ranges[i+1] - ranges[i]) + for frames in range(xpsnr_frames): + xpsnr_score = xpsnr_scores[xpsnr_iter] + xpsnr_chunk_scores.append(xpsnr_score) + xpsnr_total_scores.append(xpsnr_score) + xpsnr_iter += 1 + (xpsnr_average, xpsnr_percentile_5, xpsnr_percentile_95) = calculate_std_dev(xpsnr_chunk_scores) + xpsnr_percentile_5_total.append(xpsnr_percentile_5) + (xpsnr_average, xpsnr_percentile_5, xpsnr_percentile_95) = calculate_std_dev(xpsnr_total_scores) + + print(f'XPSNR:') + print(f'Median score: {xpsnr_average}') + print(f'5th Percentile: {xpsnr_percentile_5}') + print(f'95th Percentile: {xpsnr_percentile_95}\n') + generate_zones(ranges, xpsnr_percentile_5_total, xpsnr_average, cq, xpsnr_zones_txt_path) + + case 3: + ssimu2_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_ssimu2.log") + (ssimu2_scores, skip) = get_ssimu2(ssimu2_txt_path) + xpsnr_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_xpsnr.log") + xpsnr_scores: list[int] = get_xpsnr(xpsnr_txt_path) + + multiplied_zones_txt_path = f"{tmp_dir}/multiplied_zones.txt" + multiplied_total_scores: list[int] = [] + multiplied_percentile_5_total = [] + multiplied_iter = 0 + for i in range(len(ranges)-1): + multiplied_chunk_scores: list[int] = [] + ssimu2_frames = (ranges[i+1] - ranges[i]) // skip + for frames in range(ssimu2_frames): + ssimu2_score = ssimu2_scores[multiplied_iter] + xpsnr_index = (skip*frames) + ranges[i] + 1 + xpsnr_scores_averaged = 0 + for avg_index in range(skip): + xpsnr_scores_averaged += xpsnr_scores[xpsnr_index + avg_index - 1] + xpsnr_scores_averaged /= skip + multiplied_score = xpsnr_scores_averaged * ssimu2_score + multiplied_chunk_scores.append(multiplied_score) + multiplied_total_scores.append(multiplied_score) + multiplied_iter += 1 + (multiplied_average, multiplied_percentile_5, multiplied_percentile_95) = calculate_std_dev(multiplied_chunk_scores) + multiplied_percentile_5_total.append(multiplied_percentile_5) + (multiplied_average, multiplied_percentile_5, multiplied_percentile_95) = calculate_std_dev(multiplied_total_scores) + + print(f'Multiplied:') + print(f'Median score: {multiplied_average}') + print(f'5th Percentile: {multiplied_percentile_5}') + print(f'95th Percentile: {multiplied_percentile_95}\n') + generate_zones(ranges, multiplied_percentile_5_total, multiplied_average, cq, multiplied_zones_txt_path) + + + case 4: + ssimu2_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_ssimu2.log") + (ssimu2_scores, skip) = get_ssimu2(ssimu2_txt_path) + xpsnr_txt_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(src_file))[0]}_xpsnr.log") + xpsnr_scores: list[int] = get_xpsnr(xpsnr_txt_path) + + minimum_zones_txt_path = f"{tmp_dir}/minimum_zones.txt" + minimum_total_scores: list[int] = [] + minimum_percentile_5_total = [] + minimum_iter = 0 + ssimu2_total_scores: list[int] = [] + for ssimu2_iter in range(len(ssimu2_scores)-1): + ssimu2_total_scores.append(ssimu2_scores[ssimu2_iter]) + (ssimu2_average, ssimu2_percentile_5, ssimu2_percentile_95) = calculate_std_dev(ssimu2_total_scores) + + for i in range(len(ranges)-1): + minimum_chunk_scores: list[int] = [] + ssimu2_frames = (ranges[i+1] - ranges[i]) // skip + for frames in range(ssimu2_frames): + ssimu2_score = ssimu2_scores[minimum_iter] + xpsnr_index = (skip*frames) + ranges[i] + 1 + xpsnr_scores_averaged = 0 + for avg_index in range(skip): + xpsnr_scores_averaged += xpsnr_scores[xpsnr_index + avg_index - 1] + xpsnr_scores_averaged /= skip + xpsnr_scores_averaged *= ssimu2_average + minimum_score = min(ssimu2_score, xpsnr_scores_averaged) + minimum_chunk_scores.append(minimum_score) + minimum_total_scores.append(minimum_score) + minimum_iter += 1 + (minimum_average, minimum_percentile_5, minimum_percentile_95) = calculate_std_dev(minimum_chunk_scores) + minimum_percentile_5_total.append(minimum_percentile_5) + (minimum_average, minimum_percentile_5, minimum_percentile_95) = calculate_std_dev(minimum_total_scores) + + print(f'Minimum:') + print(f'Median score: {minimum_average}') + print(f'5th Percentile: {minimum_percentile_5}') + print(f'95th Percentile: {minimum_percentile_95}\n') + generate_zones(ranges, minimum_percentile_5_total, minimum_average, cq, minimum_zones_txt_path) + +match stage: + case 0: + workers = args.workers + crf = float(args.quality) + preset = args.preset + fast_pass(src_file, output_file, tmp_dir, preset, crf, workers) + ranges = get_ranges(scenes_file) + metrics = int(args.metrics) + calculate_metrics(src_file, output_file, tmp_dir, ranges, skip, metrics) + zones = int(args.zones) + calculate_zones(tmp_dir, ranges, zones, crf) + case 1: + workers = args.workers + crf = float(args.quality) + preset = args.preset + fast_pass(src_file, output_file, tmp_dir, preset, crf, workers) + case 2: + ranges = get_ranges(scenes_file) + metrics = int(args.metrics) + calculate_metrics(src_file, output_file, tmp_dir, ranges, skip, metrics) + case 3: + ranges = get_ranges(scenes_file) + zones = int(args.zones) + crf = float(args.quality) + calculate_zones(tmp_dir, ranges, zones, crf) + case _: + print(f"Stage argument invalid, exiting.") + exit(-2) \ No newline at end of file