diff --git a/src/TrafficLightDetector/traffic_light_detector.py b/src/TrafficLightDetector/traffic_light_detector.py
index beabd93..ff07a0c 100644
--- a/src/TrafficLightDetector/traffic_light_detector.py
+++ b/src/TrafficLightDetector/traffic_light_detector.py
@@ -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
diff --git a/src/TrafficLightDetector/traffic_light_webcam.py b/src/TrafficLightDetector/traffic_light_webcam.py
index 080dc75..0c82b9d 100644
--- a/src/TrafficLightDetector/traffic_light_webcam.py
+++ b/src/TrafficLightDetector/traffic_light_webcam.py
@@ -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