WIP fix some typing stuff

Sparks_in_Blender_encode.py

@@ -14,8 +14,10 @@ from collections import OrderedDict
 
 import json
 
+from typing import Union, Any
+
 # encoding options used
-encoding = {
+encoding: dict[str, Any] = {
     "libx264": {
         "crf": [10, 14, 16, 18, 20, 22, 25],
         "presets": [
@@ -57,8 +59,8 @@ def now():
 
 def write_line(
     codec: str,
-    crf: str,
-    preset: str,
+    crf: int,
+    preset: Union[str, int],
     infile: str,
     outfilesize: float,
     enctime: int,
@@ -70,8 +72,8 @@ def write_line(
 
     Parameters:
         codec (str): Codec used
-        crf (str): CRF used
-        preset (str): Preset used
+        crf (int): CRF used
+        preset (str/int): Preset used
         infile (str): Input file name
         outfilesize (float): Size of output file
         enctime (int): Time to encode
@@ -94,6 +96,89 @@ def write_line(
         )
 
 
+def encode_general(
+    inputfile: str, outputfile: str, codec: str, crf: int, preset: Union[str, int]
+):
+    """
+    General encoding function
+
+    Parameters:
+    inputfile (str): Path to input file
+    outputfile (str): Path to output file
+    codec (str): Codec used
+    crf (int): CRF value
+    preset (str/int): Choosen preset
+    """
+    ff = ffmpy.FFmpeg(
+        inputs={inputfile: None},
+        outputs={
+            outputfile: "-c:v {videocodec} -crf {crf} -preset {preset} -g 240 -map 0:v:0 ".format(
+                videocodec=codec,
+                crf=crf,
+                preset=preset,
+            )
+        },
+    )
+
+    return ff
+
+
+def encode_libaom(inputfile: str, outputfile: str, crf: int, preset: Union[str, int]):
+    """
+    Encoding with libaom
+
+    Parameters:
+    inputfile (str): Path to input file
+    outputfile (str): Path to output file
+    crf (int): CRF value
+    preset (str/int): Choosen preset
+    """
+    ff = ffmpy.FFmpeg(
+        inputs={inputfile: None},
+        outputs={
+            outputfile: "-c:v libaom-av1 -crf {crf} -b:v 0 -cpu-used {preset} -row-mt 1 -tiles 2x2 -g 240 -map 0:v:0 ".format(
+                crf=crf,
+                preset=preset,
+            )
+        },
+    )
+
+    return ff
+
+
+def score_vmaf(outputfile: str, inputfile: str) -> dict[str, float]:
+    """
+    Calculate a file's VMAF score
+
+    Parameters:
+    outputfile (str): Path to output file
+    inputfile (str): Path to input file
+
+    Return:
+    dict[str, float]: VMAF mean and min value
+    """
+    ff = ffmpy.FFmpeg(
+        inputs=OrderedDict([(outputfile, None), (inputfile, None)]),
+        outputs={
+            "-": "-filter_complex libvmaf=log_fmt=json:n_threads={cputhreads}:log_path=vmaf.json -f null".format(
+                cputhreads=cpu_count()
+            )
+        },
+    )
+
+    ff.run()
+
+    with open("vmaf.json", "r") as file:
+        vmafall = json.load(file)
+
+    vmaf: dict[str, float] = {
+        "mean": vmafall["pooled_metrics"]["vmaf"]["mean"],
+        "min": vmafall["pooled_metrics"]["vmaf"]["min"],
+    }
+
+    return vmaf
+
+
 if __name__ == "__main__":
     if not os.path.isdir("encoded"):
         os.mkdir("encoded")
@@ -115,6 +200,8 @@ if __name__ == "__main__":
             )
         )
 
+    inputfile = "Sparks_in_Blender.webm"
+
     for codec in encoding:
         for crf in encoding[codec]["crf"]:
             for preset in encoding[codec]["presets"]:
@@ -133,28 +220,22 @@ if __name__ == "__main__":
 
                 # libaom needs additional options
                 if codec == "libaom-av1":
-                    ff = ffmpy.FFmpeg(
-                        inputs={"Sparks_in_Blender.webm": None},
-                        outputs={
-                            outputfile: "-c:v {videocodec} -crf {crf} -b:v 0 -cpu-used {preset} -row-mt 1 -tiles 2x2 -g 240 -map 0:v:0 ".format(
-                                videocodec=codec,
-                                crf=crf,
-                                preset=preset,
-                            )
-                        },
+                    ff = encode_libaom(
+                        inputfile=inputfile,
+                        outputfile=outputfile,
+                        crf=crf,
+                        preset=preset,
                     )
                 else:
-                    ff = ffmpy.FFmpeg(
-                        inputs={"Sparks_in_Blender.webm": None},
-                        outputs={
-                            outputfile: "-c:v {videocodec} -crf {crf} -preset {preset} -g 240 -map 0:v:0 ".format(
-                                videocodec=codec,
-                                crf=crf,
-                                preset=preset,
-                            )
-                        },
+                    ff = encode_general(
+                        inputfile=inputfile,
+                        outputfile=outputfile,
+                        codec=codec,
+                        crf=crf,
+                        preset=preset,
                     )
 
+                # execute previously defined encoding settings
                 starttime = now()
                 ff.run()
                 endtime = now()
@@ -162,21 +243,7 @@ if __name__ == "__main__":
 
                 outputfilesize = os.path.getsize(outputfile) / 1024 / 1024
 
-                ffvmaf = ffmpy.FFmpeg(
-                    inputs=OrderedDict(
-                        [(outputfile, None), ("Sparks_in_Blender.webm", None)]
-                    ),
-                    outputs={
-                        "-": "-filter_complex libvmaf=log_fmt=json:n_threads={cputhreads}:log_path=vmaf.json -f null".format(
-                            cputhreads=cpu_count()
-                        )
-                    },
-                )
-
-                ffvmaf.run()
-
-                with open("vmaf.json", "r") as file:
-                    vmaf = json.load(file)
+                vmaf = score_vmaf(outputfile=outputfile, inputfile=inputfile)
 
                 write_line(
                     codec=codec,
@@ -185,8 +252,8 @@ if __name__ == "__main__":
                     infile="Sparks_in_Blender.webm",
                     outfilesize=outputfilesize,
                     enctime=difftime,
-                    vmafmean=vmaf["pooled_metrics"]["vmaf"]["mean"],
-                    vmafmin=vmaf["pooled_metrics"]["vmaf"]["min"],
+                    vmafmean=vmaf["mean"],
+                    vmafmin=vmaf["min"],
                 )
 
                 os.remove(outputfile)