kristoferssolo/Traffic-Light-Detector
1
0
Fork 0
mirror of https://github.com/kristoferssolo/Traffic-Light-Detector.git synced 2025-10-21 20:00:36 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Created class ColorRange

Browse Source
This commit is contained in:
Kristofers Solo 2022-12-24 11:36:54 +02:00
parent 455226432a
commit fac36ed9bd
2 changed files with 25 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
src/TrafficLightDetector/color.py
View File

@ -1,14 +1,19 @@
from typing import NamedTuple
import cv2 import cv2
import numpy as np import numpy as np
from loguru import logger from loguru import logger
class ColorRange(NamedTuple):
color1: tuple[int, int, int]
color2: tuple[int, int, int]
@logger.catch @logger.catch
class Color: class ColorAttributes:
def __init__(self, name: str, color: tuple[int, int, int], def __init__(self, name: str, color: tuple[int, int, int], lowers: ColorRange,
lowers: tuple[tuple[int, int, int], tuple[int, int, int]], uppers: ColorRange, hsv: cv2.cvtColor, minDist: int, param2: int) -> None:
uppers: tuple[tuple[int, int, int], tuple[int, int, int]],
hsv: cv2.cvtColor, minDist: int, param2: int) -> None:
self.name = name self.name = name
self.color = color self.color = color
@ -22,6 +27,7 @@ class Color:
self.mask = cv2.add(self.mask, mask) self.mask = cv2.add(self.mask, mask)
# set circle # set circle
self.circle = cv2.HoughCircles(self.mask, cv2.HOUGH_GRADIENT, 1, minDist=minDist, param1=50, param2=param2, minRadius=0, maxRadius=30) self.circle = cv2.HoughCircles(self.mask, cv2.HOUGH_GRADIENT, 1, minDist=minDist, param1=50,
param2=param2, minRadius=0, maxRadius=30)
if self.circle is not None: if self.circle is not None:
self.circle = np.uint16(np.around(self.circle)) self.circle = np.uint16(np.around(self.circle))

25
src/TrafficLightDetector/traffic_light_detector.py
View File

@ -1,6 +1,7 @@
import cv2 import cv2
from paths import HAAR_PATH
from TrafficLightDetector.color import Color from TrafficLightDetector.color import ColorAttributes, ColorRange
from TrafficLightDetector.paths import HAAR_PATH
class TrafficLightDetector: class TrafficLightDetector:
@ -10,12 +11,12 @@ class TrafficLightDetector:
BOUNDARY = 2 BOUNDARY = 2
TEXT = False TEXT = False
# HSV values # HSV values
RED_LOWER = ((160, 100, 100), (0, 100, 100)) RED_LOWER = ColorRange((160, 100, 100), (0, 100, 100))
RED_UPPER = ((180, 255, 255), (10, 255, 255)) RED_UPPER = ColorRange((180, 255, 255), (10, 255, 255))
YELLOW_LOWER = ((15, 150, 150), (15, 100, 100)) YELLOW_LOWER = ColorRange((15, 150, 150), (15, 100, 100))
YELLOW_UPPER = ((35, 255, 255), (35, 255, 255)) YELLOW_UPPER = ColorRange((35, 255, 255), (35, 255, 255))
GREEN_LOWER = ((40, 50, 50), (95, 45, 38)) GREEN_LOWER = ColorRange((40, 50, 50), (95, 45, 38))
GREEN_UPPER = ((90, 255, 255), (130, 60, 60)) GREEN_UPPER = ColorRange((90, 255, 255), (130, 60, 60))
# BGR values # BGR values
RED = (0, 0, 200) RED = (0, 0, 200)
@ -25,11 +26,11 @@ class TrafficLightDetector:
def _set_image(self, image=None, roi=None, detectTrafficLights=True) -> None: def _set_image(self, image=None, roi=None, detectTrafficLights=True) -> None:
self.image = image self.image = image
self.roi = self.image if roi is None else roi self.roi = self.image if roi is None else roi
self.size = self.image.shape if roi is None else self.roi.shape self.size = self.roi.shape if roi is not None else self.image.shape
hsv = cv2.cvtColor(self.image if roi is None else self.roi, cv2.COLOR_BGR2HSV) 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.red = ColorAttributes("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.yellow = ColorAttributes("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.green = ColorAttributes("GREEN", self.GREEN, self.GREEN_LOWER, self.GREEN_UPPER, hsv, minDist=30, param2=5)
self.colors = [self.red, self.yellow, self.green] self.colors = [self.red, self.yellow, self.green]
self.detect_traffic_lights = detectTrafficLights self.detect_traffic_lights = detectTrafficLights
self.signal_color = "" self.signal_color = ""