Combined TrafficLight and color recognition

src/TrafficLightDetector/traffic_light_detector.py

@@ -1,4 +1,5 @@
 import cv2
+import numpy as np
 from loguru import logger
 from paths import HAAR_PATH
 from TrafficLightDetector.color import Color
@@ -22,9 +23,9 @@ class TrafficLightDetector:
     YELLOW = (0, 175, 225)
     GREEN = (0, 150, 0)
 
-    def _set_image(self, image=None) -> None:
+    def _set_image(self, image) -> None:
         self.image = image
-        self.image_copy = self.image
+        self.roi_color = self.image
         self.size = self.image.shape
         hsv = cv2.cvtColor(self.image, cv2.COLOR_BGR2HSV)
         self.red = Color("RED", self.RED, self.RED_LOWER, self.RED_UPPER, hsv, minDist=80, param2=10)
@@ -32,13 +33,25 @@ class TrafficLightDetector:
         self.green = Color("GREEN", self.GREEN, self.GREEN_LOWER, self.GREEN_UPPER, hsv, minDist=30, param2=5)
         self.colors = [self.red, self.yellow, self.green]
 
-    def _find_traffic_lights(self) -> None:
+    def _set_roi(self) -> None:
+        self.size = self.roi_color.shape
+        hsv = cv2.cvtColor(self.roi_color, cv2.COLOR_BGR2HSV)
+        self.red = Color("RED", self.RED, self.RED_LOWER, self.RED_UPPER, hsv, minDist=80, param2=10)
+        self.yellow = Color("YELLOW", self.YELLOW, self.YELLOW_LOWER, self.YELLOW_UPPER, hsv, minDist=60, param2=10)
+        self.green = Color("GREEN", self.GREEN, self.GREEN_LOWER, self.GREEN_UPPER, hsv, minDist=30, param2=5)
+        self.colors = [self.red, self.yellow, self.green]
+
+    def _outline_traffic_lights(self, image) -> None:
+        self.image = image
         gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
         # draw rectangle around traffic lights
         for x, y, width, height in self.CASCADE.detectMultiScale(gray, 1.2, 5):
             cv2.rectangle(self.image, (x, y), (x + width, y + height), (255, 0, 0), self.BOUNDARY)
+            self.roi_color = self.image[y:y + height, x:x + width]
+            self._set_roi()
+            self._draw_circle()
 
-    def _draw_circle(self) -> None:
+    def _draw_circle(self, text=False) -> None:
         for color in self.colors:
             if color.circle is not None:
                 for values in color.circle[0, :]:
@@ -53,9 +66,10 @@ class TrafficLightDetector:
                             h += color.mask[values[1] + inner_radius, values[0] + outter_radius]
                             s += 1
                     if h / s > 100:
-                        cv2.circle(self.image_copy, (values[0], values[1]), values[2] + 10, color.color, 2)
+                        cv2.circle(self.roi_color, (values[0], values[1]), values[2] + 10, color.color, 2)  # draws circle
                         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)
+                        if text:
+                            cv2.putText(self.roi_color, color.name, (values[0], values[1]), self.FONT, 1, color.color, 2, cv2.LINE_AA)  # draws text
                         self.signal = color.name
 
     def get_signal(self) -> str:

src/TrafficLightDetector/traffic_light_webcam.py

@@ -1,4 +1,5 @@
 import cv2
+from loguru import logger
 from TrafficLightDetector.traffic_light_detector import TrafficLightDetector
 
 
@@ -11,11 +12,8 @@ class TrafficLightDetectorWebcam(TrafficLightDetector):
         while True:
             _, frame = self.video_capture.read()
 
-            # self._set_image(frame)
+            self._outline_traffic_lights(frame)
-            # self._draw_circle()
+            cv2.imshow("Video", self.image)
-            self._find_traffic_lights()
-            cv2.imshow("Video", self.image_copy)
-            cv2.imshow("Video", frame)
             if cv2.waitKey(1) & 0xFF == ord("q"):
                 break
         self.video_capture.release()