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