From b359ec7ed18ac0f141c0adb8feee55433fc4d6b8 Mon Sep 17 00:00:00 2001
From: Kristofers Solo <kristians.cagulis@gmail.com>
Date: Wed, 21 Dec 2022 13:25:21 +0200
Subject: [PATCH] Added detection, when signal is green

---
 main.py                                       | 26 ++++++++++++++-----
 .../traffic_light_detector.py                 |  7 +++--
 .../traffic_light_webcam.py                   | 23 ++++++++++++----
 3 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/main.py b/main.py
index 1b2f84d..f0d922f 100755
--- a/main.py
+++ b/main.py
@@ -6,6 +6,12 @@ from paths import create_dirs, IMAGES_IN_PATH
 
 
 parser = argparse.ArgumentParser(description="Traffic light detection script.")
+parser.add_argument(
+    "-i",
+    "--image",
+    action="store_true",
+    help="detects traffic lights in images located in ./assets/images_in/",
+)
 parser.add_argument(
     "-w",
     "--webcam",
@@ -13,20 +19,16 @@ parser.add_argument(
     help="reads webcam inputs to determine traffic light color",
 )
 parser.add_argument(
-    "-i",
-    "--image",
+    "-c",
+    "--change",
     action="store_true",
-    help="detects traffic lights in images located in ./assets/images_in/",
+    help="detects traffic lights change",
 )
 
 
 @logger.catch
 def main(args) -> None:
     create_dirs()
-    if args.webcam:
-        from TrafficLightDetector.traffic_light_webcam import TrafficLightDetectorWebcam
-        camera = TrafficLightDetectorWebcam(0)  # Change number if webcam didn't detect
-        camera.enable()
 
     if args.image:
         from TrafficLightDetector.traffic_light_images import TrafficLightDetectorImages
@@ -34,6 +36,16 @@ def main(args) -> None:
             image = TrafficLightDetectorImages(path)
             image.draw()
 
+    if args.webcam:
+        from TrafficLightDetector.traffic_light_webcam import TrafficLightDetectorWebcam
+        camera = TrafficLightDetectorWebcam(0)  # Change number if webcam didn't detect
+        camera.enable()
+
+    if args.change:
+        from TrafficLightDetector.traffic_light_webcam import TrafficLightDetectorWebcam
+        camera = TrafficLightDetectorWebcam(0, sound=True)  # Change number if webcam didn't detect
+        camera.enable()
+
 
 if __name__ == "__main__":
     args = parser.parse_args()
diff --git a/src/TrafficLightDetector/traffic_light_detector.py b/src/TrafficLightDetector/traffic_light_detector.py
index 5cd9e5a..8ad99ff 100644
--- a/src/TrafficLightDetector/traffic_light_detector.py
+++ b/src/TrafficLightDetector/traffic_light_detector.py
@@ -22,6 +22,8 @@ class TrafficLightDetector:
     YELLOW = (0, 175, 225)
     GREEN = (0, 150, 0)
 
+    signal_color = ""
+
     def _set_image(self, image=None, roi=None, detectTrafficLights=True) -> None:
         self.image = image
         self.roi = self.image if roi is None else roi
@@ -62,7 +64,4 @@ class TrafficLightDetector:
                         if self.TEXT:
                             cv2.putText(self.roi if self.detect_traffic_lights else self.image, color.name,
                                         (value[0], value[1]), self.FONT, 1, color.color, 2, cv2.LINE_AA)  # draws text
-                        self.signal = color.name
-
-    def get_signal(self) -> str:
-        return self.signal
+                        self.signal_color = color.name
diff --git a/src/TrafficLightDetector/traffic_light_webcam.py b/src/TrafficLightDetector/traffic_light_webcam.py
index 694154c..8d83d47 100644
--- a/src/TrafficLightDetector/traffic_light_webcam.py
+++ b/src/TrafficLightDetector/traffic_light_webcam.py
@@ -1,19 +1,32 @@
 import cv2
 from TrafficLightDetector.traffic_light_detector import TrafficLightDetector
+from loguru import logger
 
 
 class TrafficLightDetectorWebcam(TrafficLightDetector):
 
-    def __init__(self, source: int) -> None:
+    def __init__(self, source: int, sound: bool = False) -> None:
         self.video_capture = cv2.VideoCapture(source)
+        self.sound = sound
 
     def enable(self) -> None:
         while True:
-            _, frame = self.video_capture.read()
-            self._set_image(frame)
-            self._outline_traffic_lights()
-            cv2.imshow("Video", self.image)
+            self._get_video()
+            if self.sound:
+                self._make_sound()
             if cv2.waitKey(1) & 0xFF == ord("q"):
                 break
         self.video_capture.release()
         cv2.destroyAllWindows()
+
+    def _get_video(self) -> None:
+        _, frame = self.video_capture.read()
+        self._set_image(frame)
+        self._outline_traffic_lights()
+        cv2.imshow("Video", self.image)
+
+    def _make_sound(self) -> None:
+        """Do some sound if green light"""
+        if self.signal_color == "GREEN":
+            # DO SOME SOUND
+            logger.debug("DRIVE")