1#!/usr/bin/env python3 2 3# E.g. `./gen_combos.py [--write] color_quads/720p.png` 4 5import concurrent.futures 6import pathlib 7import subprocess 8import sys 9 10ARGS = sys.argv 11SRC_PATH = pathlib.Path(ARGS.pop()) 12assert SRC_PATH.exists(), "gen_combos.py [--flags] <src file path>" 13DIR = SRC_PATH.parent 14 15 16# crossCombine([{a:false},{a:5}], [{},{b:5}]) 17# [{a:false}, {a:true}, {a:false,b:5}, {a:true,b:5}] 18def cross_combine(*args): 19 args = list(args) 20 21 def cross_combine2(listA, listB): 22 listC = [] 23 for a in listA: 24 for b in listB: 25 c = dict() 26 c.update(a) 27 c.update(b) 28 listC.append(c) 29 return listC 30 31 res = [dict()] 32 while True: 33 try: 34 next = args.pop(0) 35 except IndexError: 36 break 37 res = cross_combine2(res, next) 38 return res 39 40 41def keyed_combiner(key, vals): 42 res = [] 43 for v in vals: 44 d = dict() 45 d[key] = v 46 res.append(d) 47 return res 48 49 50# - 51 52 53def eprint(*args, **kwargs): 54 print(*args, file=sys.stderr, **kwargs) 55 56 57# - 58 59OGG = [] 60WEBM_CODECS = ["av1", "vp9"] 61 62if "--all" in ARGS: 63 OGG = cross_combine( 64 [{"ext": "ogg"}], keyed_combiner("vcodec", ["theora", "vp8", "vp9"]) 65 ) 66 WEBM_CODECS += ["vp8"] 67 68MP4 = cross_combine([{"ext": "mp4"}], keyed_combiner("vcodec", ["av1", "h264", "vp9"])) 69 70WEBM = cross_combine([{"ext": "webm"}], keyed_combiner("vcodec", WEBM_CODECS)) 71 72# - 73 74FORMAT_LIST = set( 75 [ 76 "yuv420p", 77 "yuv420p10", 78 # 'yuv420p12', 79 # 'yuv420p16be', 80 # 'yuv420p16le', 81 "gbrp", 82 ] 83) 84 85if "--all" in ARGS: 86 FORMAT_LIST |= set( 87 [ 88 "yuv420p", 89 "yuv420p10", 90 "yuv420p12", 91 "yuv420p16be", 92 "yuv420p16le", 93 "yuv422p", 94 "yuv422p10", 95 "yuv422p12", 96 "yuv422p16be", 97 "yuv422p16le", 98 "yuv444p", 99 "yuv444p10", 100 "yuv444p12", 101 "yuv444p16be", 102 "yuv444p16le", 103 "yuv411p", 104 "yuv410p", 105 "yuyv422", 106 "uyvy422", 107 "rgb24", 108 "bgr24", 109 "rgb8", 110 "bgr8", 111 "rgb444be", 112 "rgb444le", 113 "bgr444be", 114 "bgr444le", 115 # 'nv12', # Encoding not different than yuv420p? 116 # 'nv21', # Encoding not different than yuv420p? 117 "gbrp", 118 "gbrp9be", 119 "gbrp9le", 120 "gbrp10be", 121 "gbrp10le", 122 "gbrp12be", 123 "gbrp12le", 124 "gbrp14be", 125 "gbrp14le", 126 "gbrp16be", 127 "gbrp16le", 128 ] 129 ) 130 131FORMATS = keyed_combiner("format", list(FORMAT_LIST)) 132 133RANGE = keyed_combiner("range", ["tv", "pc"]) 134 135CSPACE_LIST = set( 136 [ 137 "bt709", 138 # 'bt2020', 139 ] 140) 141 142if "--all" in ARGS: 143 CSPACE_LIST |= set( 144 [ 145 "bt709", 146 "bt2020", 147 "bt601-6-525", # aka smpte170m NTSC 148 "bt601-6-625", # aka bt470bg PAL 149 ] 150 ) 151CSPACE_LIST = list(CSPACE_LIST) 152 153# - 154 155COMBOS = cross_combine( 156 WEBM + MP4 + OGG, 157 FORMATS, 158 RANGE, 159 keyed_combiner("src_cspace", CSPACE_LIST), 160 keyed_combiner("dst_cspace", CSPACE_LIST), 161) 162 163# - 164 165print(f"{len(COMBOS)} combinations...") 166 167todo = [] 168for c in COMBOS: 169 dst_name = ".".join( 170 [ 171 SRC_PATH.name, 172 c["src_cspace"], 173 c["dst_cspace"], 174 c["range"], 175 c["format"], 176 c["vcodec"], 177 c["ext"], 178 ] 179 ) 180 181 src_cspace = c["src_cspace"] 182 183 vf = f"scale=out_range={c['range']}" 184 vf += f",colorspace=all={c['dst_cspace']}" 185 vf += f":iall={src_cspace}" 186 args = [ 187 "ffmpeg", 188 "-y", 189 # For input: 190 "-color_primaries", 191 src_cspace, 192 "-color_trc", 193 src_cspace, 194 "-colorspace", 195 src_cspace, 196 "-i", 197 SRC_PATH.as_posix(), 198 # For output: 199 "-bitexact", # E.g. don't use true random uuids 200 "-vf", 201 vf, 202 "-pix_fmt", 203 c["format"], 204 "-vcodec", 205 c["vcodec"], 206 "-crf", 207 "1", # Not-quite-lossless 208 (DIR / dst_name).as_posix(), 209 ] 210 if "-v" in ARGS or "-vv" in ARGS: 211 print("$ " + " ".join(args)) 212 else: 213 print(" " + args[-1]) 214 215 todo.append(args) 216 217# - 218 219with open(DIR / "reftest.list", "r") as f: 220 reftest_list_text = f.read() 221 222for args in todo: 223 vid_name = pathlib.Path(args[-1]).name 224 if vid_name not in reftest_list_text: 225 print(f"WARNING: Not in reftest.list: {vid_name}") 226 227# - 228 229if "--write" not in ARGS: 230 print("Use --write to write. Exiting...") 231 exit(0) 232 233# - 234 235 236def run_cmd(args): 237 dest = None 238 if "-vv" not in ARGS: 239 dest = subprocess.DEVNULL 240 try: 241 subprocess.run(args, stderr=dest) 242 except FileNotFoundError: 243 print("FileNotFoundError, is ffmpeg not in your PATH?") 244 raise 245 246 247with concurrent.futures.ThreadPoolExecutor() as pool: 248 fs = [] 249 for cur_args in todo: 250 f = pool.submit(run_cmd, cur_args) 251 fs.append(f) 252 253 done = 0 254 for f in concurrent.futures.as_completed(fs): 255 f.result() # Raise if it raised 256 done += 1 257 sys.stdout.write(f"\rEncoded {done}/{len(todo)}") 258