109 lines
4.0 KiB
Python
Executable File
109 lines
4.0 KiB
Python
Executable File
import statistics
|
|
from math import ceil
|
|
import json
|
|
import sys
|
|
import subprocess
|
|
import vapoursynth as vs
|
|
core = vs.core
|
|
|
|
if "--help" in sys.argv[1:]:
|
|
print('Usage:\npython auto-boost_2.0.py "{animu.mkv}" {base CQ/CRF/Q}"\n\nExample:\npython "auto-boost_2.0.py" "path/to/nice_boat.mkv" 30')
|
|
exit(0)
|
|
else:
|
|
pass
|
|
|
|
og_cq = int(sys.argv[2]) # CQ to start from
|
|
br = 10 # maximum CQ change from original
|
|
|
|
def get_ranges(scenes):
|
|
ranges = []
|
|
ranges.insert(0,0)
|
|
with open(scenes, "r") as file:
|
|
content = json.load(file)
|
|
for i in range(len(content['scenes'])):
|
|
ranges.append(content['scenes'][i]['end_frame'])
|
|
return ranges
|
|
|
|
iter = 0
|
|
def zones_txt(beginning_frame, end_frame, cq, zones_loc):
|
|
global iter
|
|
iter += 1
|
|
|
|
with open(zones_loc, "w" if iter == 1 else "a") as file:
|
|
file.write(f"{beginning_frame} {end_frame} svt-av1 --crf {cq}\n")
|
|
|
|
def calculate_standard_deviation(score_list: list[int]):
|
|
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)
|
|
return (average, sorted_score_list[len(filtered_score_list)//20])
|
|
|
|
fast_av1an_command = f'av1an -i "{sys.argv[1]}" --temp "{sys.argv[1][:-4]}/temp/" -y \
|
|
--verbose --keep --split-method av-scenechange -m lsmash \
|
|
--min-scene-len 12 -c mkvmerge --sc-downscale-height 480 \
|
|
--set-thread-affinity 2 -e svt-av1 --force -v \" \
|
|
--preset 9 --crf {og_cq} --rc 0 --film-grain 0 --lp 2 \
|
|
--scm 0 --keyint 0 --fast-decode 1 --color-primaries 1 \
|
|
--transfer-characteristics 1 --matrix-coefficients 1 \" \
|
|
--pix-format yuv420p10le -x 240 -w {WORKERS} \
|
|
-o "{sys.argv[1][:-4]}_fastpass.mkv"'
|
|
|
|
p = subprocess.Popen(fast_av1an_command, shell=True)
|
|
exit_code = p.wait()
|
|
|
|
if exit_code != 0:
|
|
print("Av1an encountered an error, exiting.")
|
|
exit(-2)
|
|
|
|
scenes_loc = f"{sys.argv[1][:-4]}/temp/scenes.json"
|
|
ranges = get_ranges(scenes_loc)
|
|
|
|
src = core.lsmas.LWLibavSource(source=sys.argv[1], cache=0)
|
|
enc = core.lsmas.LWLibavSource(source=f"{sys.argv[1][:-4]}_fastpass.mkv", cache=0)
|
|
|
|
print(f"source: {len(src)} frames")
|
|
print(f"encode: {len(enc)} frames")
|
|
|
|
source_clip = src.resize.Bicubic(format=vs.RGBS, matrix_in_s='709').fmtc.transfer(transs="srgb", transd="linear", bits=32)
|
|
encoded_clip = enc.resize.Bicubic(format=vs.RGBS, matrix_in_s='709').fmtc.transfer(transs="srgb", transd="linear", bits=32)
|
|
|
|
percentile_5_total = []
|
|
total_ssim_scores: list[int] = []
|
|
|
|
skip = 10 # amount of skipped frames
|
|
|
|
for i in range(len(ranges)-1):
|
|
cut_source_clip = source_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=0)
|
|
cut_encoded_clip = encoded_clip[ranges[i]:ranges[i+1]].std.SelectEvery(cycle=skip, offsets=0)
|
|
result = cut_source_clip.ssimulacra2.SSIMULACRA2(cut_encoded_clip)
|
|
chunk_ssim_scores: list[int] = []
|
|
|
|
for index, frame in enumerate(result.frames()):
|
|
score = frame.props['_SSIMULACRA2']
|
|
# print(f'Frame {index}/{result.num_frames}: {score}')
|
|
chunk_ssim_scores.append(score)
|
|
total_ssim_scores.append(score)
|
|
|
|
(average, percentile_5) = calculate_standard_deviation(chunk_ssim_scores)
|
|
percentile_5_total.append(percentile_5)
|
|
|
|
(average, percentile_5) = calculate_standard_deviation(total_ssim_scores)
|
|
print(f'Median score: {average}\n\n')
|
|
|
|
for i in range(len(ranges)-1):
|
|
|
|
new_cq = og_cq - ceil((1.0 - (percentile_5_total[i]/average)) / 0.5 * 10) # trust me bro
|
|
|
|
if new_cq < og_cq-br: # set lowest allowed cq
|
|
new_cq = og_cq-br
|
|
|
|
if new_cq > og_cq+br: # set highest allowed cq
|
|
new_cq = og_cq+br
|
|
|
|
print(f'Enc: [{ranges[i]}:{ranges[i+1]}]\n'
|
|
f'Chunk 5th percentile: {percentile_5_total[i]}\n'
|
|
f'Adjusted CRF: {new_cq}\n\n')
|
|
zones_txt(ranges[i], ranges[i+1], new_cq, f"{scenes_loc[:-11]}zones.txt")
|
|
|
|
# yes, this is messier than the 1.0 code, laziness won over me, deal with it
|