From ed8ee4259303652a2f8a3860427538937eeda165 Mon Sep 17 00:00:00 2001
From: trix <93526043+trixoniisama@users.noreply.github.com>
Date: Wed, 3 Jul 2024 00:41:53 +0200
Subject: [PATCH] Add auto-boost algorithm 2.0 + fixes

---
 README.md         |  48 ++++++++++++++++++++-
 auto-boost_2.0.py | 106 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 152 insertions(+), 2 deletions(-)
 create mode 100755 auto-boost_2.0.py

diff --git a/README.md b/README.md
index e050b5a..7b3fe2a 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,9 @@
-# Auto-boost Algorithms
+# Auto-boost Algorithm
 
 ## 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.
+The requirements are Vapoursynth, vstools and LSMASHSource.
 
 __Usage:__
 ```
@@ -14,3 +14,47 @@ __Example:__
 ```
 python auto-boost_1.0.py "path/to/nice_boat.mkv" "path/to/scenes.json" 30
 ```
+
+__Advantages:__
+- Fast
+- No bs
+- Solves one long-lasting issue of AV1 encoders: low bitrate allocation in dark scenes
+
+__Known limitations:__
+- Not every dark scene is made equal, brightness is not a great enough metric to determine whether CRF should be decreased or not
+- CRF is boosted to the max during credits
+- Script now entirely irrelevant with SVT-AV1-PSY's new frame-luma-bias feature
+
+_Inspiration was drawn from the original Av1an (python) boosting code_
+
+## Version 2.0: SSIMULACRA2-based
+
+> Does a fast encode of the provided file, calculates SSIMULACRA2 scores of each chunks and adjusts CRF per-scene to be closer to the average total score, in a _zones.txt_ file to feed Av1an.
+The requirements are Vapoursynth, vstools, LSMASHSource, Av1an and vapoursynth-ssimulacra2.
+
+__Usage:__
+```
+python auto-boost_2.0.py "{animu.mkv}" "{scenes.json}" {base CQ/CRF/Q}
+```
+
+__Example:__
+```
+python auto-boost_2.0.py "path/to/nice_boat.mkv" "path/to/scenes.json" 30
+```
+
+__Advantages:__
+- Lower quality deviation of individual scenes in regards to the entire stream
+- Better allocates bitrate in more complex scenes and compensates by giving less bitrate to scenes presenting some headroom for further compression
+
+__Known limitations:__
+- Slow process
+- No bitrate cap in place so the size of complex scenes can go out of hand
+- The SSIMULACRA2 metric is not ideal, plus the score alone is not representative enough of if a CRF adjustement is relevant in the context of that scene (AI will save)
+
+_Borrowed some code from Sav1or's SSIMULACRA2 script_
+
+## Version 3.0: SSIMULACRA2-based + per-scene grain synthesis strength determination
+
+...and a few other improvements.
+
+_Soon:tm:_
diff --git a/auto-boost_2.0.py b/auto-boost_2.0.py
new file mode 100755
index 0000000..782315a
--- /dev/null
+++ b/auto-boost_2.0.py
@@ -0,0 +1,106 @@
+import statistics
+from math import ceil
+import json
+import sys
+from vapoursynth import vs, core
+
+if "--help" in sys.argv[1:]:
+    print('Usage:\npython auto-boost_2.0.py "{animu.mkv}" "{scenes.json}" {base CQ/CRF/Q}"\n\nExample:\npython "auto-boost_2.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
+
+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])
+
+scenes_loc = sys.argv[2] # scene file is expected to be named 'scenes.json'
+ranges = get_ranges(scenes_loc)
+
+fast_av1an_command = f'av1an -i "{sys.argv[1]}" --temp "{scenes_loc[:-11]}/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)
+
+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