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