diff --git a/README.md b/README.md index ae80c6f..e050b5a 100644 --- a/README.md +++ b/README.md @@ -1 +1,16 @@ -# auto-boost-algorithm \ No newline at end of file +# Auto-boost Algorithms + +## Version 1.0: brightness-based + +> Gets the average brightness of a scene and lowers CQ/CRF/Q the darker the scene is, in a _zones.txt_ file to feed Av1an. +Only requirements are Vapoursynth and vstools. + +__Usage:__ +``` +python auto-boost_1.0.py "{animu.mkv}" "{scenes.json}" {base CQ/CRF/Q} "{encoder: aom/svt-av1/rav1e (optional)}" +``` + +__Example:__ +``` +python auto-boost_1.0.py "path/to/nice_boat.mkv" "path/to/scenes.json" 30 +``` diff --git a/auto-boost_1.0.py b/auto-boost_1.0.py new file mode 100755 index 0000000..3c42b11 --- /dev/null +++ b/auto-boost_1.0.py @@ -0,0 +1,98 @@ +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])