99 lines
3.1 KiB
Python
Executable File
99 lines
3.1 KiB
Python
Executable File
import statistics
|
|
from math import ceil
|
|
import json
|
|
import sys
|
|
from vapoursynth import core
|
|
from vstools import clip_async_render
|
|
|
|
if "--help" in sys.argv[1:]:
|
|
print('Usage:\npython auto-boost_1.0.py "{animu.mkv}" "{scenes.json}" {base CQ/CRF/Q} "{encoder: aom/svt-av1/rav1e (optional)}"\n\nExample:\npython "auto-boost_1.0.py" "path/to/nice_boat.mkv" "path/to/scenes.json" 30')
|
|
exit(0)
|
|
else:
|
|
pass
|
|
|
|
og_cq = int(sys.argv[3]) # CQ to start from
|
|
br = 10 # maximum CQ change from original
|
|
bl = og_cq-br # hard cap how low CQ can be set by boost
|
|
|
|
try:
|
|
ENCODER = sys.argv[4] # select encoder between aom, svt-av1 and rav1e :allhailav1:
|
|
except:
|
|
ENCODER = "svt-av1"
|
|
|
|
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
|
|
|
|
def get_brightness(video, start, end):
|
|
brightness = []
|
|
ref = video[:].std.PlaneStats(plane=0)
|
|
|
|
render = clip_async_render(
|
|
ref, outfile=None, progress=f'Getting frame props... from {start} to {end}',
|
|
callback=lambda _, f: f.props.copy()
|
|
)
|
|
props = [prop['PlaneStatsAverage'] for prop in render]
|
|
|
|
for prop in props:
|
|
brightness.append(prop)
|
|
|
|
brig_geom = round(statistics.geometric_mean([x+0.01 for x in brightness]), 2) #x+1
|
|
#print(brig_geom)
|
|
|
|
return brig_geom
|
|
|
|
def boost(br_geom):
|
|
global br, bl
|
|
global og_cq
|
|
|
|
if br_geom < 0.5: # too dark, cq needs to change
|
|
new_cq = og_cq - ceil((0.5 - br_geom) / 0.5 * br)
|
|
|
|
if new_cq < bl: # Cap on boosting
|
|
new_cq = bl
|
|
|
|
return new_cq
|
|
|
|
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:
|
|
if ENCODER == "aom":
|
|
file.write(f"{beginning_frame} {end_frame} aom --cq-level={cq}\n")
|
|
elif ENCODER == "svt-av1":
|
|
file.write(f"{beginning_frame} {end_frame} svt-av1 --crf {cq}\n")
|
|
elif ENCODER == "rav1e":
|
|
file.write(f"{beginning_frame} {end_frame} rav1e --quantizer {cq}\n")
|
|
else:
|
|
print("Incompatible encoder given.")
|
|
exit(-1)
|
|
|
|
def zones_main(chunk, start, end, zones_loc):
|
|
global og_cq
|
|
|
|
br = get_brightness(chunk, start, end) # brightness range is [0,1]
|
|
|
|
cq = boost(br)
|
|
|
|
if og_cq != cq and cq != None:
|
|
print(f'Enc: [{start}:{end}]\n'
|
|
f'Avg brightness: {br}\n'
|
|
f'Adjusted CQ: {cq}\n\n')
|
|
zones_txt(start, end, cq, f"{zones_loc}zones.txt")
|
|
elif cq == None:
|
|
print(f"cq = None (brightness > 0.5, actually: {br})")
|
|
|
|
scenes_loc = sys.argv[2] # scene file is expected to be named 'scenes.json'
|
|
ranges = get_ranges(scenes_loc)
|
|
src = core.lsmas.LWLibavSource(source=sys.argv[1], cache=0)
|
|
for i in range(len(ranges)-1):
|
|
#print(f"[{ranges[i]}:{ranges[i+1]}]")
|
|
zones_main(src[ranges[i]:ranges[i+1]], ranges[i], ranges[i+1], scenes_loc[:-11])
|