diff --git a/src/TrafficLightDetector/traffic_light_detector.py b/src/TrafficLightDetector/traffic_light_detector.py index 58d787b..bf6257e 100644 --- a/src/TrafficLightDetector/traffic_light_detector.py +++ b/src/TrafficLightDetector/traffic_light_detector.py @@ -1,5 +1,4 @@ import cv2 -import numpy as np from loguru import logger from paths import HAAR_PATH from TrafficLightDetector.color import Color @@ -10,6 +9,7 @@ class TrafficLightDetector: FONT = cv2.FONT_HERSHEY_SIMPLEX RADIUS = 5 BOUNDARY = 2 + TEXT = False # HSV values RED_LOWER = ((160, 100, 100), (0, 100, 100)) RED_UPPER = ((180, 255, 255), (10, 255, 255)) @@ -23,53 +23,46 @@ class TrafficLightDetector: YELLOW = (0, 175, 225) GREEN = (0, 150, 0) - def _set_image(self, image) -> None: + def _set_image(self, image=None, roi=None, detectTrafficLights=True) -> None: self.image = image - self.roi_color = self.image - self.size = self.image.shape - hsv = cv2.cvtColor(self.image, cv2.COLOR_BGR2HSV) + self.roi = self.image if roi is None else roi + self.size = self.image.shape if roi is None else self.roi.shape + hsv = cv2.cvtColor(self.image if roi is None else self.roi, 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] + self.detect_traffic_lights = detectTrafficLights - 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 + def _outline_traffic_lights(self) -> None: 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() + roi = self.image[y:y + height, x:x + width] + self._set_image(self.image, roi) self._draw_circle() - def _draw_circle(self, text=False) -> None: + def _draw_circle(self) -> None: for color in self.colors: if color.circle is not None: - 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: + for value in color.circle[0, :]: + if value[0] > self.size[1] or value[1] > self.size[0] or value[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 (values[1] + inner_radius) >= self.size[0] or (values[0] + outter_radius) >= self.size[1]: + for i in range(-self.RADIUS, self.RADIUS): + for j in range(-self.RADIUS, self.RADIUS): + if (value[1] + i) >= self.size[0] or (value[0] + j) >= self.size[1]: continue - h += color.mask[values[1] + inner_radius, values[0] + outter_radius] + h += color.mask[value[1] + i, value[0] + j] s += 1 if h / s > 100: - 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) - if text: - cv2.putText(self.roi_color, color.name, (values[0], values[1]), self.FONT, 1, color.color, 2, cv2.LINE_AA) # draws text + cv2.circle(self.roi if self.detect_traffic_lights else self.image, (value[0], value[1]), value[2] + 10, color.color, 2) # draws circle + cv2.circle(color.mask, (value[0], value[1]), value[2] + 30, (255, 255, 255), 2) + 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: diff --git a/src/TrafficLightDetector/traffic_light_webcam.py b/src/TrafficLightDetector/traffic_light_webcam.py index 0ff5fe2..6b729d4 100644 --- a/src/TrafficLightDetector/traffic_light_webcam.py +++ b/src/TrafficLightDetector/traffic_light_webcam.py @@ -11,8 +11,8 @@ class TrafficLightDetectorWebcam(TrafficLightDetector): def enable(self) -> None: while True: _, frame = self.video_capture.read() - - self._outline_traffic_lights(frame) + self._set_image(frame) + self._outline_traffic_lights() cv2.imshow("Video", self.image) if cv2.waitKey(1) & 0xFF == ord("q"): break