diff --git a/code/perception/src/traffic_light_node.py b/code/perception/src/traffic_light_node.py index 00034bab..66b250ee 100755 --- a/code/perception/src/traffic_light_node.py +++ b/code/perception/src/traffic_light_node.py @@ -11,6 +11,8 @@ from cv_bridge import CvBridge from traffic_light_detection.src.traffic_light_detection.traffic_light_inference \ import TrafficLightInference # noqa: E501 +import cv2 +import numpy as np class TrafficLightNode(CompatibleNode): @@ -58,12 +60,17 @@ def auto_invalidate_state(self): self.last_info_time = None def handle_camera_image(self, image): - result, data = self.classifier(self.bridge.imgmsg_to_cv2(image)) + cv2_image = self.bridge.imgmsg_to_cv2(image) + rgb_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB) + result, data = self.classifier(cv2_image) if data[0][0] > 1e-15 and data[0][3] > 1e-15 or \ data[0][0] > 1e-10 or data[0][3] > 1e-10: return # too uncertain, may not be a traffic light + if not is_front(rgb_image): + return # not a front facing traffic light + state = result if result in [1, 2, 4] else 0 if self.last_state == state: # 1: Green, 2: Red, 4: Yellow, 0: Unknown @@ -81,6 +88,46 @@ def run(self): self.spin() +def get_light_mask(image): + hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + + # Define lower and upper bounds for the hue, saturation, and value + # - Red and Yellow combined since they are so close in the color spectrum + lower_red_yellow = np.array([0, 75, 100]) + upper_red_yellow = np.array([40, 255, 255]) + lower_green = np.array([40, 200, 200]) + upper_green = np.array([80, 255, 255]) + + # Mask where the pixels within the bounds are white, otherwise black + m1 = cv2.inRange(hsv, lower_red_yellow, upper_red_yellow) + m2 = cv2.inRange(hsv, lower_green, upper_green) + mask = cv2.bitwise_or(m1, m2) + + return mask + + +def is_front(image): + mask = get_light_mask(image) + + # Find contours in the thresholded image, use only the largest one + contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, + cv2.CHAIN_APPROX_SIMPLE) + contours = sorted(contours, key=cv2.contourArea, reverse=True)[:1] + contour = contours[0] if contours else None + + if contour is None: + return False + + _, _, width, height = cv2.boundingRect(contour) + aspect_ratio = width / height + + # If aspect ratio is within range of a square (therefore a circle) + if 0.75 <= aspect_ratio <= 1.3: + return True + else: + return False + + if __name__ == "__main__": roscomp.init("TrafficLightNode") node = TrafficLightNode("TrafficLightNode")