From a2d16d4dc47bed3158323cc5b8a67edf58475970 Mon Sep 17 00:00:00 2001 From: Kristofers Solo Date: Sat, 17 Dec 2022 16:11:49 +0200 Subject: [PATCH] Moved to folder --- pyproject.toml | 2 +- src/{ => TrafficLightDetector}/__init__.py | 0 src/TrafficLightDetector/color.py | 27 ++++++++ src/{ => TrafficLightDetector}/py.typed | 0 .../traffic_light_detector.py | 54 ++++++++++++++++ .../traffic_light_images.py | 20 ++++++ .../traffic_light_webcam.py | 20 ++++++ src/detect.py | 62 ------------------- src/paths.py | 2 +- 9 files changed, 123 insertions(+), 64 deletions(-) rename src/{ => TrafficLightDetector}/__init__.py (100%) create mode 100644 src/TrafficLightDetector/color.py rename src/{ => TrafficLightDetector}/py.typed (100%) create mode 100644 src/TrafficLightDetector/traffic_light_detector.py create mode 100644 src/TrafficLightDetector/traffic_light_images.py create mode 100644 src/TrafficLightDetector/traffic_light_webcam.py delete mode 100644 src/detect.py diff --git a/pyproject.toml b/pyproject.toml index 3bd5311..ef66097 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=42.0", "wheel"] build-backend = "setuptools.build_meta" [tool.pytest.ini_options] -addops = "--cov=detector" +addops = "--cov=TrafficLightDetector" testpaths = ["tests"] [tool.mypy] diff --git a/src/__init__.py b/src/TrafficLightDetector/__init__.py similarity index 100% rename from src/__init__.py rename to src/TrafficLightDetector/__init__.py diff --git a/src/TrafficLightDetector/color.py b/src/TrafficLightDetector/color.py new file mode 100644 index 0000000..f058b9d --- /dev/null +++ b/src/TrafficLightDetector/color.py @@ -0,0 +1,27 @@ +import cv2 +import numpy as np +from loguru import logger + + +@logger.catch +class Color: + def __init__(self, name: str, color: tuple[int, int, int], + lowers: tuple[tuple[int, int, int], tuple[int, int, int]], + uppers: tuple[tuple[int, int, int], tuple[int, int, int]], + hsv: cv2.cvtColor, minDist: int, param2: int) -> None: + self.name = name + self.color = color + + # set mask + masks: list[cv2.inRange] = [] + for lower, upper in zip(lowers, uppers): + masks.append(cv2.inRange(hsv, lower, upper)) + self.mask = masks[0] + if len(masks) > 1: + for mask in masks: + self.mask = cv2.add(self.mask, mask) + + # set circle + 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: + self.circle = np.uint16(np.around(self.circle)) diff --git a/src/py.typed b/src/TrafficLightDetector/py.typed similarity index 100% rename from src/py.typed rename to src/TrafficLightDetector/py.typed diff --git a/src/TrafficLightDetector/traffic_light_detector.py b/src/TrafficLightDetector/traffic_light_detector.py new file mode 100644 index 0000000..4ddc6ad --- /dev/null +++ b/src/TrafficLightDetector/traffic_light_detector.py @@ -0,0 +1,54 @@ +import cv2 +from loguru import logger +from TrafficLightDetector.color import Color + + +class TrafficLightDetector: + FONT = cv2.FONT_HERSHEY_SIMPLEX + RADIUS = 5 + BOUNDARY = 4 / 10 + # HSV values + RED_LOWER = ((160, 100, 100), (0, 100, 100)) + RED_UPPER = ((180, 255, 255), (10, 255, 255)) + YELLOW_LOWER = ((15, 150, 150), (15, 100, 100)) + YELLOW_UPPER = ((35, 255, 255), (35, 255, 255)) + GREEN_LOWER = ((40, 50, 50), (95, 45, 38)) + GREEN_UPPER = ((90, 255, 255), (130, 60, 60)) + + # BGR values + RED = (0, 0, 200) + YELLOW = (0, 175, 225) + GREEN = (0, 150, 0) + + def _set_image(self, image) -> None: + self.image = image + self.image_copy = 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) + 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 _draw_circle(self): + try: + 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: + 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]: + continue + h += color.mask[i[1] + inner_radius, i[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) + except AttributeError: + logger.warning("Image/frame was not specified") diff --git a/src/TrafficLightDetector/traffic_light_images.py b/src/TrafficLightDetector/traffic_light_images.py new file mode 100644 index 0000000..64163d7 --- /dev/null +++ b/src/TrafficLightDetector/traffic_light_images.py @@ -0,0 +1,20 @@ +from pathlib import Path + +import cv2 +from paths import IMAGES_OUT_PATH +from TrafficLightDetector.traffic_light_detector import TrafficLightDetector + + +class TrafficLightDetectorImages(TrafficLightDetector): + + def __init__(self, path) -> None: + self.path = path + self.image = cv2.imread(str(path)) + super().__init__(self.image) + + def _save_image(self) -> None: + cv2.imwrite(str(IMAGES_OUT_PATH.joinpath(self.path.name)), self.image_copy) + + def draw(self) -> None: + self._draw_circle(self.image) + self._save_image() diff --git a/src/TrafficLightDetector/traffic_light_webcam.py b/src/TrafficLightDetector/traffic_light_webcam.py new file mode 100644 index 0000000..1cf1ee1 --- /dev/null +++ b/src/TrafficLightDetector/traffic_light_webcam.py @@ -0,0 +1,20 @@ +import cv2 +from loguru import logger +from TrafficLightDetector.traffic_light_detector import TrafficLightDetector + + +class TrafficLightDetectorWebcam(TrafficLightDetector): + + def __init__(self) -> None: + self.video_capture = cv2.VideoCapture(0) # Change number if webcam didn't detect + + def enable(self) -> None: + while True: + _, image = self.video_capture.read() + self._set_image(image) + self._draw_circle() + cv2.imshow("Video", self.image_copy) + if cv2.waitKey(1) & 0xFF == ord("q"): + break + self.video_capture.release() + cv2.destroyAllWindows() diff --git a/src/detect.py b/src/detect.py deleted file mode 100644 index 7d015c8..0000000 --- a/src/detect.py +++ /dev/null @@ -1,62 +0,0 @@ -import cv2 -import numpy as np -from loguru import logger -from paths import IMAGES_OUT_PATH -from pathlib import Path - - -@logger.catch -class Color: - - def __init__(self, name: str, lower: tuple[int, int, int], upper: tuple[int, int, int], hsv: cv2.cvtColor, minDist: int, param2: int) -> None: - self.name: str = name - self.lower: tuple[int, int, int] = lower - self.upper: tuple[int, int, int] = upper - self.mask = cv2.inRange(hsv, self.lower, self.upper) - 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: - self.circle = np.uint16(np.around(self.circle)) - - -@logger.catch -class TrafficLightDetector: - FONT = cv2.FONT_HERSHEY_SIMPLEX - RADIUS = 5 - BOUNDARY = 4 / 10 - - def __init__(self, path: Path) -> None: - self.path = path - self.image = cv2.imread(str(self.path)) - self.image_copy = self.image - self.size = self.image.shape - # self.red1 = Color("RED", (0, 100, 100), (10, 255, 255)) - hsv = cv2.cvtColor(self.image, cv2.COLOR_BGR2HSV) - self.red = Color("RED", (160, 100, 100), (180, 255, 255), hsv, minDist=80, param2=10) - self.yellow = Color("YELLOW", (15, 150, 150), (35, 255, 255), hsv, minDist=60, param2=10) - self.green = Color("GREEN", (40, 50, 50), (90, 255, 255), hsv, minDist=30, param2=5) - self.colors = [self.red, self.yellow, self.green] - - def draw_circle(self) -> None: - for color in self.colors: - if color.circle is not None: - - for i in color.circle[0, :]: - logger.debug(f"{i = }") - if i[0] > self.size[1] or i[1] > self.size[0] or i[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]: - continue - h += color.mask[i[1] + inner_radius, i[0] + outter_radius] - s += 1 - if h / s > 100: - cv2.circle(self.image_copy, (i[0], i[1]), i[2] + 10, (0, 255, 0), 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, (255, 0, 0), 2, cv2.LINE_AA) - self._save_image() - - def _save_image(self) -> None: - cv2.imwrite(str(IMAGES_OUT_PATH.joinpath(self.path.name)), self.image_copy) diff --git a/src/paths.py b/src/paths.py index ae27aba..37da27d 100644 --- a/src/paths.py +++ b/src/paths.py @@ -15,7 +15,7 @@ PATHS = (LOGS_PATH, IMAGES_IN_PATH, IMAGES_OUT_PATH) log_level = "DEBUG" if BASE_PATH.joinpath("debug").exists() else "INFO" # Set up logging -logger.add(LOGS_PATH.joinpath("detection.log"), format="{time} | {level} | {message}", level=log_level, rotation="1 MB", compression="zip") +logger.add(LOGS_PATH.joinpath("detection.log"), format="{time} | {level} | {message}", level=log_level, rotation="10 MB", compression="zip") @logger.catch