Add auto-boost algorithm 1.0

README.md

@@ -1 +1,16 @@
-# auto-boost-algorithm
+# 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
+```

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])