From 070a130558815faa9bb126913f9dcccdc60d826f Mon Sep 17 00:00:00 2001
From: RealStickman <mrc@frm01.net>
Date: Thu, 10 Nov 2022 13:38:11 +0100
Subject: [PATCH] Add argparse

---
 encode_single_video.py | 62 +++++++++++++++++++++++++++++++++++-------
 1 file changed, 52 insertions(+), 10 deletions(-)

diff --git a/encode_single_video.py b/encode_single_video.py
index e2a6872..e38a8a6 100755
--- a/encode_single_video.py
+++ b/encode_single_video.py
@@ -8,6 +8,8 @@ import csv
 
 import ffmpy
 
+import argparse
+
 from multiprocessing import cpu_count
 
 from collections import OrderedDict
@@ -364,17 +366,57 @@ def main(inputfile: str, outputpath: str = "encodes"):
                     mse=mse,
                 )
 
-                # TODO make removal optional/togglable with cli switch
-                os.remove(outputfile)
-
-                os.remove("vmaf.json")
-                os.remove("ssim.log")
-                os.remove("psnr.log")
+                cleanup(keepencodes=keepencodes, outputfile=outputfile)
 
 
 if __name__ == "__main__":
-    # TODO expose this as variable/argument?
-    outputpath: str = ""
-    inputfile: str = "source/Sparks_in_Blender.webm"
+    parser = argparse.ArgumentParser(description="")
 
-    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)