191 - Filter out side-facing traffic lights by HSV and shape

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