Added TrafficLight recognition

main.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import argparse
-from loguru import logger
 
+from loguru import logger
 from paths import create_dirs, IMAGES_IN_PATH
 
 

src/TrafficLightDetector/traffic_light_detector.py

@@ -35,21 +35,25 @@ class TrafficLightDetector:
             for color in self.colors:
                 if color.circle is not None:
 
-                    for i in color.circle[0, :]:
-                        if i[0] > self.size[1] or i[1] > self.size[0] or i[1] > self.size[0] * self.BOUNDARY:
+                    logger.debug(f"{color.circle = }")
+                    for values in color.circle[0, :]:
+                        if values[0] > self.size[1] or values[1] > self.size[0] or values[1] > self.size[0] * self.BOUNDARY:
                             continue
 
                         h, s = 0, 0
                         for inner_radius in range(-self.RADIUS, self.RADIUS):
                             for outter_radius in range(-self.RADIUS, self.RADIUS):
-                                if (i[1] + inner_radius) >= self.size[0] or (i[0] + outter_radius) >= self.size[1]:
+                                if (values[1] + inner_radius) >= self.size[0] or (values[0] + outter_radius) >= self.size[1]:
                                     continue
-                                h += color.mask[i[1] + inner_radius, i[0] + outter_radius]
+                                h += color.mask[values[1] + inner_radius, values[0] + outter_radius]
                                 s += 1
                         if h / s > 100:
-                            cv2.circle(self.image_copy, (i[0], i[1]), i[2] + 10, color.color, 2)
-                            cv2.circle(color.mask, (i[0], i[1]), i[2] + 30, (255, 255, 255), 2)
-                            cv2.putText(self.image_copy, color.name, (i[0], i[1]), self.FONT, 1, color.color, 2, cv2.LINE_AA)
-                            logger.debug(color.name)
+                            cv2.circle(self.image_copy, (values[0], values[1]), values[2] + 10, color.color, 2)
+                            cv2.circle(color.mask, (values[0], values[1]), values[2] + 30, (255, 255, 255), 2)
+                            cv2.putText(self.image_copy, color.name, (values[0], values[1]), self.FONT, 1, color.color, 2, cv2.LINE_AA)
+                            self.signal = color.name
         except AttributeError:
             logger.warning("Image/frame was not specified")
+
+    def get_signal(self) -> str:
+        return self.signal

src/TrafficLightDetector/traffic_light_webcam.py

@@ -1,5 +1,6 @@
 import cv2
 from loguru import logger
+from paths import HAAR_PATH
 from TrafficLightDetector.traffic_light_detector import TrafficLightDetector
 
 
@@ -7,13 +8,21 @@ class TrafficLightDetectorWebcam(TrafficLightDetector):
 
     def __init__(self) -> None:
         self.video_capture = cv2.VideoCapture(0)  # Change number if webcam didn't detect
+        self.lights_cascade = cv2.CascadeClassifier(str(HAAR_PATH))
 
     def enable(self) -> None:
         while True:
-            _, image = self.video_capture.read()
-            self._set_image(image)
-            self._draw_circle()
-            cv2.imshow("Video", self.image_copy)
+            _, frame = self.video_capture.read()
+            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+            lights = self.lights_cascade.detectMultiScale(gray, 1.2, 5)
+            for x, y, w, h in lights:
+                cv2.rectangle(frame, (x, y), (x+w, y+h), (255, 0, 0), 5)
+
+            # self._set_image(frame)
+            # self._draw_circle()
+            # cv2.imshow("Video", self.image_copy)
+            cv2.imshow("Video", frame)
             if cv2.waitKey(1) & 0xFF == ord("q"):
                 break
         self.video_capture.release()

src/paths.py

@@ -9,6 +9,7 @@ LOGS_PATH = BASE_PATH.joinpath(".logs")
 ASSETS_PATH = BASE_PATH.joinpath("assets")
 IMAGES_IN_PATH = ASSETS_PATH.joinpath("images_in")
 IMAGES_OUT_PATH = ASSETS_PATH.joinpath("images_out")
+HAAR_PATH = ASSETS_PATH.joinpath("haar").joinpath("TrafficLights.xml")
 
 PATHS = (LOGS_PATH, IMAGES_IN_PATH, IMAGES_OUT_PATH)