From 939501896cbc790f0d2cb9acbc77192bf728f132 Mon Sep 17 00:00:00 2001
From: Kristofers Solo <kristians.cagulis@gmail.com>
Date: Sun, 11 Dec 2022 17:46:26 +0200
Subject: [PATCH] Added debug file

---
 .gitignore |  1 +
 main.py    | 20 ++++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 1376d7f..06c1951 100644
--- a/.gitignore
+++ b/.gitignore
@@ -130,3 +130,4 @@ dmypy.json
 
 assets/*
 model.h5
+debug
diff --git a/main.py b/main.py
index 92ff1a3..3ef02a0 100755
--- a/main.py
+++ b/main.py
@@ -1,8 +1,16 @@
 #!/usr/bin/env python3
 import argparse
 import importlib
+from pathlib import Path
 
-from detector.paths import create_dirs
+from detector.paths import BASE_PATH, create_dirs, LOGS_PATH
+from loguru import logger
+
+log_level = "DEBUG" if Path.exists(Path.joinpath(BASE_PATH, "debug")) else "INFO"
+
+
+# Set up logging
+logger.add(Path.joinpath(LOGS_PATH, "detection.log"), format="{time} | {level} | {message}", level=log_level, rotation="1 MB", compression="zip")
 
 
 parser = argparse.ArgumentParser(description="Traffic light detection script.")
@@ -30,8 +38,13 @@ parser.add_argument(
     action="store_true",
     help="detects traffic lights in videos located in ./assets/detection/videos_in/",
 )
+parser.add_argument(
+    "--test",
+    action="store_true",
+)
 
 
+@logger.catch
 def main(args) -> None:
     create_dirs()
     if args.extract:
@@ -45,7 +58,10 @@ def main(args) -> None:
         module.detect_traffic_light_color_image()
     if args.video:
         module = importlib.import_module("detector.detect_traffic_light_color_video")
-        module.detect_traffic_light_color_image()
+        module.detect_traffic_light_color_video()
+    if args.test:
+        module = importlib.import_module("detector.test")
+        module.main()
 
 
 if __name__ == "__main__":