Restructuring to function both as cli and import

encode_single_video.py

@@ -58,6 +58,7 @@ def now():
 
 
 def write_line(
+    datafile: str,
     codec: str,
     crf: int,
     preset: Union[str, int],
@@ -73,6 +74,7 @@ def write_line(
     Write line to data csv
 
     Parameters:
+    datafile (str): Path to output CSV filename
     codec (str): Codec used
     crf (int): CRF used
     preset (str/int): Preset used
@@ -241,7 +243,7 @@ def score_psnr(outputfile: str, inputfile: str) -> float:
     inputfile (str): Path to input file
 
     Return:
-    TBD
+    TODO
     """
     ff = ffmpy.FFmpeg(
         inputs=OrderedDict([(outputfile, None), (inputfile, None)]),
@@ -265,12 +267,23 @@ def score_psnr(outputfile: str, inputfile: str) -> float:
     return mse_avg
 
 
-if __name__ == "__main__":
-    if not os.path.isdir("encodes"):
-        os.mkdir("encodes")
+def main(inputfile: str, outputpath: str = "encodes"):
+    """
+    Main program function so this program can be used from the command line or as a library import.
+
+    Parameters:
+    inputfile (str): Path to input video file
+
+    Optional Parameters:
+    outputpath (str): Path to output folder. Defaults to "encodes" in the current working directory
+    """
+    if not os.path.isdir(outputpath):
+        os.mkdir(outputpath)
 
     datafile = version + "-data-" + str(now()) + ".csv"
 
+    inputfilename = os.path.basename(os.path.splitext(inputfile)[0])
+
     with open(datafile, "w", newline="") as file:
         write = csv.writer(file)
         write.writerow(
@@ -288,17 +301,15 @@ if __name__ == "__main__":
             )
         )
 
-    # TODO expose this as variable/argument?
-    inputfile = "source/Sparks_in_Blender.webm"
-
     for codec in encoding:
         for crf in encoding[codec]["crf"]:
             for preset in encoding[codec]["presets"]:
                 # TODO selection of output location with arguments?
                 outputfile = os.path.join(
-                    "encodes",
+                    outputpath,
                     (
-                        "Sparks_in_Blender-codec_"
+                        inputfilename
+                        + "-codec_"
                         + codec
                         + "-crf_"
                         + str(crf)
@@ -340,10 +351,11 @@ if __name__ == "__main__":
                 mse = score_psnr(outputfile=outputfile, inputfile=inputfile)
 
                 write_line(
+                    datafile=datafile,
                     codec=codec,
                     crf=crf,
                     preset=preset,
-                    infile="Sparks_in_Blender.webm",
+                    infile=os.path.basename(inputfile),
                     outfilesize=outputfilesize,
                     enctime=difftime,
                     vmafmean=vmaf["mean"],
@@ -358,3 +370,11 @@ if __name__ == "__main__":
                 os.remove("vmaf.json")
                 os.remove("ssim.log")
                 os.remove("psnr.log")
+
+
+if __name__ == "__main__":
+    # TODO expose this as variable/argument?
+    outputpath: str = ""
+    inputfile: str = "source/Sparks_in_Blender.webm"
+
+    main(inputfile=inputfile)