From f4b3a42466b0d23a1ae9054d95258769d2946737 Mon Sep 17 00:00:00 2001 From: Kristofers Solo Date: Sat, 17 Dec 2022 16:40:07 +0200 Subject: [PATCH] Added TrafficLight recognition --- main.py | 2 +- .../traffic_light_detector.py | 20 +++++++++++-------- .../traffic_light_webcam.py | 17 ++++++++++++---- src/paths.py | 1 + 4 files changed, 27 insertions(+), 13 deletions(-) diff --git a/main.py b/main.py index 5c2b6bd..6a31acb 100755 --- a/main.py +++ b/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 diff --git a/src/TrafficLightDetector/traffic_light_detector.py b/src/TrafficLightDetector/traffic_light_detector.py index 4a32739..beabd93 100644 --- a/src/TrafficLightDetector/traffic_light_detector.py +++ b/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 diff --git a/src/TrafficLightDetector/traffic_light_webcam.py b/src/TrafficLightDetector/traffic_light_webcam.py index 1cf1ee1..080dc75 100644 --- a/src/TrafficLightDetector/traffic_light_webcam.py +++ b/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() diff --git a/src/paths.py b/src/paths.py index 37da27d..3145405 100644 --- a/src/paths.py +++ b/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)