kristoferssolo/Traffic-Light-Detector
1
0
Fork 0
mirror of https://github.com/kristoferssolo/Traffic-Light-Detector.git synced 2026-03-22 00:36:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added Traffic Light recognition to class

Browse Source
This commit is contained in:
Kristofers Solo
2022-12-17 16:51:13 +02:00
parent 0ee644320d
commit dced61a204
2 changed files with 29 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
src/TrafficLightDetector/traffic_light_detector.py
View File

@@ -1,12 +1,14 @@
import cv2
from loguru import logger
from paths import HAAR_PATH
from TrafficLightDetector.color import Color
class TrafficLightDetector:
CASCADE = cv2.CascadeClassifier(str(HAAR_PATH))
FONT = cv2.FONT_HERSHEY_SIMPLEX
RADIUS = 5
BOUNDARY = 4 / 10
BOUNDARY = 2
# HSV values
RED_LOWER = ((160, 100, 100), (0, 100, 100))
RED_UPPER = ((180, 255, 255), (10, 255, 255))
@@ -20,7 +22,7 @@ class TrafficLightDetector:
YELLOW = (0, 175, 225)
GREEN = (0, 150, 0)
def _set_image(self, image) -> None:
def _set_image(self, image=None) -> None:
self.image = image
self.image_copy = self.image
self.size = self.image.shape
@@ -30,30 +32,31 @@ 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:
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)
def _draw_circle(self) -> None:
try:
for color in self.colors:
if color.circle is not 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:
continue
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 (values[1] + inner_radius) >= self.size[0] or (values[0] + outter_radius) >= self.size[1]:
continue
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(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")
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]:
continue
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(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
def get_signal(self) -> str:
return self.signal

11
src/TrafficLightDetector/traffic_light_webcam.py
View File

@@ -1,6 +1,4 @@
import cv2
from loguru import logger
from paths import HAAR_PATH
from TrafficLightDetector.traffic_light_detector import TrafficLightDetector
@@ -8,20 +6,15 @@ 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:
_, 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)
self._find_traffic_lights()
cv2.imshow("Video", self.image_copy)
cv2.imshow("Video", frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break