Updated argument handling

main.py

@@ -1,43 +1,53 @@
-#!/bin/python3.10
+#!/usr/bin/env python3
-import sys
+import argparse
+import importlib
 
 from detector.paths import create_dirs
 
-ARGS = """
-Usage: main.py [OPTION]
 
-Options:
+parser = argparse.ArgumentParser(description="Traffic light detection script.")
-    -h --help       Displays this list
+parser.add_argument(
-    -e --extract    Excracts and cropps traffic light images from given images in ./assets/exctraction/input/ to ./assets/exctraction/cropped/
+    "-e",
-    -t --train      Trains model
+    "--extract",
-    -i --image      Detecs traffic lights in images located in ./assets/detection/images_in/
+    action="store_true",
-    -v --video      Detecs traffic lights in videos located in ./assets/detection/videos_in/
+    help="extracts and crops traffic light images from given images in ./assets/extraction/input/ to ./assets/extraction/cropped/",
-"""
+)
+parser.add_argument(
+    "-t",
+    "--train",
+    action="store_true",
+    help="trains model.",
+)
+parser.add_argument(
+    "-i",
+    "--image",
+    action="store_true",
+    help="detects traffic lights in images located in ./assets/detection/images_in/",
+)
+parser.add_argument(
+    "-v",
+    "--video",
+    action="store_true",
+    help="detects traffic lights in videos located in ./assets/detection/videos_in/",
+)
 
 
-def main(argv) -> None:
+def main(args) -> None:
     create_dirs()
-    for arg in argv:
+    if args.extract:
-        if arg in ("-h", "--help"):
+        module = importlib.import_module("detector.extract_traffic_lights")
-            print(ARGS)
+        module.extract_traffic_lights()
-            sys.exit()
+    if args.train:
-        elif arg in ("-e", "--extract"):
+        module = importlib.import_module("detector.train_traffic_light_color")
-            from detector.extract_traffic_lights import extract_traffic_lights
+        module.train_traffic_light_color()
-            extract_traffic_lights()
+    if args.image:
-        elif arg in ("-t", "--train"):
+        module = importlib.import_module("detector.detect_traffic_light_color_image")
-            from detector.train_traffic_light_color import train_traffic_light_color
+        module.detect_traffic_light_color_image()
-            train_traffic_light_color()
+    if args.video:
-        elif arg in ("-i", "--image"):
+        module = importlib.import_module("detector.detect_traffic_light_color_video")
-            from detector.detect_traffic_light_color_image import (
+        module.detect_traffic_light_color_image()
-                detect_traffic_light_color_image,
-            )
-            detect_traffic_light_color_image()
-        elif arg in ("-v", "--video"):
-            from detector.detect_traffic_light_color_video import (
-                detect_traffic_light_color_video,
-            )
-            detect_traffic_light_color_video()
 
 
 if __name__ == "__main__":
-    main(sys.argv[1:])
+    args = parser.parse_args()
+    main(args)