Moved to folder

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]

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))

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")

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()

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()

src/detect.py

@@ -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)

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