Cleanup function on keyboard interrupt or exit

encode_single_video.py

@@ -8,6 +8,8 @@ import csv
 
 import ffmpy
 
+import argparse
+
 from multiprocessing import cpu_count
 
 from collections import OrderedDict
@@ -57,6 +59,34 @@ def now():
     return int(time.time())
 
 
+def cleanup(keepencodes: bool, outputfile: str = ""):
+    """
+    Clean up files on program exit/interruption
+    """
+
+    try:
+        # REVIEW make removal optional/togglable with cli switch
+        if not keepencodes:
+            os.remove(outputfile)
+    except:
+        pass
+
+    try:
+        os.remove("vmaf.json")
+    except:
+        pass
+
+    try:
+        os.remove("ssim.log")
+    except:
+        pass
+
+    try:
+        os.remove("psnr.log")
+    except:
+        pass
+
+
 def write_line(
     datafile: str,
     codec: str,
@@ -267,7 +297,7 @@ def score_psnr(outputfile: str, inputfile: str) -> float:
     return mse_avg
 
 
-def main(inputfile: str, outputpath: str = "encodes"):
+def main(inputfile: str, outputpath: str = "encodes", keepencodes: bool = False):
     """
     Main program function so this program can be used from the command line or as a library import.
 
@@ -364,17 +394,57 @@ def main(inputfile: str, outputpath: str = "encodes"):
                     mse=mse,
                 )
 
-                # TODO make removal optional/togglable with cli switch
+                cleanup(keepencodes=keepencodes, outputfile=outputfile)
-                os.remove(outputfile)
-
-                os.remove("vmaf.json")
-                os.remove("ssim.log")
-                os.remove("psnr.log")
 
 
 if __name__ == "__main__":
-    # TODO expose this as variable/argument?
+    parser = argparse.ArgumentParser(description="")
-    outputpath: str = ""
-    inputfile: str = "source/Sparks_in_Blender.webm"
 
-    main(inputfile=inputfile)
+    parser.add_argument(
+        "-o",
+        "--output-dir",
+        required=False,
+        type=str,
+        help='Output directory for encodes. Default is "encodes" in current working directory',
+        default="encodes",
+    )
+
+    parser.add_argument(
+        "-i",
+        "--input-file",
+        required=True,
+        type=str,
+        help="Input file to encode",
+    )
+
+    # in case you wanted to rerun the conversion for everything
+    parser.add_argument(
+        "-r",
+        "--reset",
+        required=False,
+        action="store_true",
+        help="Rerun conversion for all files",
+    )
+
+    parser.add_argument(
+        "-k",
+        "--keep-encodes",
+        action="store_true",
+        required=False,
+        help="Don't delete encodes after getting their scores. Default is delete (False)",
+    )
+
+    args = parser.parse_args()
+
+    # REVIEW expose this as variable/argument?
+    outputpath: str = args.output_dir
+    inputfile: str = args.input_file
+
+    keepencodes: bool = args.keep_encodes
+
+    # main program loop with cleanup in case of interrupt
+    try:
+        main(inputfile=inputfile, outputpath=outputpath, keepencodes=keepencodes)
+    except KeyboardInterrupt:
+        # run cleanup function on keyboard interrupt
+        cleanup(keepencodes=keepencodes)