From a8220d9a6c900dc0273f8ea645f37f2a34727583 Mon Sep 17 00:00:00 2001
From: Kristofers Solo <kristians.cagulis@gmail.com>
Date: Sat, 17 Dec 2022 16:12:18 +0200
Subject: [PATCH] Updated main

---
 main.py | 49 ++++++++++++++-----------------------------------
 1 file changed, 14 insertions(+), 35 deletions(-)

diff --git a/main.py b/main.py
index 4de7bf4..5c2b6bd 100755
--- a/main.py
+++ b/main.py
@@ -1,59 +1,38 @@
 #!/usr/bin/env python3
 import argparse
-import importlib
-
-from detector.paths import BASE_PATH, create_dirs, LOGS_PATH
 from loguru import logger
 
-log_level = "DEBUG" if BASE_PATH.joinpath("debug").exists() else "INFO"
-
-
-# Set up logging
-logger.add(LOGS_PATH.joinpath("detection.log"), format="{time} | {level} | {message}", level=log_level, rotation="1 MB", compression="zip")
+from paths import create_dirs, IMAGES_IN_PATH
 
 
 parser = argparse.ArgumentParser(description="Traffic light detection script.")
 parser.add_argument(
-    "-e",
-    "--extract",
+    "-w",
+    "--webcam",
     action="store_true",
-    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.",
+    help="reads webcam inputs to determine traffic light color",
 )
 parser.add_argument(
     "-i",
     "--image",
     action="store_true",
-    help="detects traffic lights in images located in ./assets/detection/images_in/",
+    help="detects traffic lights in images located in ./assets/images_in/",
 )
-# parser.add_argument(
-#     "-v",
-#     "--video",
-#     action="store_true",
-#     help="detects traffic lights in videos located in ./assets/detection/videos_in/",
-# )
 
 
 @logger.catch
 def main(args) -> None:
     create_dirs()
-    if args.extract:
-        module = importlib.import_module("detector.extract_traffic_lights")
-        module.extract_traffic_lights()
-    if args.train:
-        module = importlib.import_module("detector.train_traffic_light_color")
-        module.train_traffic_light_color()
+    if args.webcam:
+        from TrafficLightDetector.traffic_light_webcam import TrafficLightDetectorWebcam
+        camera = TrafficLightDetectorWebcam()
+        camera.enable()
+
     if args.image:
-        module = importlib.import_module("detector.detect_traffic_light_color_image")
-        module.detect_traffic_light_color_image()
-    # if args.video:
-    #     module = importlib.import_module("detector.detect_traffic_light_color_video")
-    #     module.detect_traffic_light_color_video()
+        from TrafficLightDetector.traffic_light_images import TrafficLightDetectorImages
+        for path in IMAGES_IN_PATH.iterdir():
+            image = TrafficLightDetectorImages(path)
+            image.draw()
 
 
 if __name__ == "__main__":