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