tl_detector.py

#!/usr/bin/env python
import rospy
from timeit import default_timer
from std_msgs.msg import Int32
from geometry_msgs.msg import PoseStamped, Pose
from styx_msgs.msg import TrafficLightArray, TrafficLight
from styx_msgs.msg import Lane
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from light_classification.tl_classifier import TLClassifier
import tf
import cv2
import yaml

from scipy.spatial import KDTree
import numpy as np

STATE_COUNT_THRESHOLD = 3
CAMERA_IMAGE_PROCESSING_RATE = 0.250 # process camera image every 250 milliseconds
# don't process camera image if waypoints between car and stop line waypoint are 
# below waypoint threshold 
WAYPOINT_THRESHOLD = 150

class TLDetector(object):
    def __init__(self):
        rospy.init_node('tl_detector')

        self.pose = None
        self.waypoints = None
        self.waypoints_2d = None
        self.waypoint_tree = None
        self.camera_image = None
        self.last_image_processed = 0 # set to 0 since we haven't processed an image yet
        self.lights = []

        sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
        sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)

        '''
        /vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
        helps you acquire an accurate ground truth data source for the traffic light
        classifier by sending the current color state of all traffic lights in the
        simulator. When testing on the vehicle, the color state will not be available. You'll need to
        rely on the position of the light and the camera image to predict it.
        '''
        sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray, self.traffic_cb)
        sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)

        config_string = rospy.get_param("/traffic_light_config")
        self.config = yaml.load(config_string)

        self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint', Int32, queue_size=1)

        self.bridge = CvBridge()
        self.light_classifier = TLClassifier(self.config['is_site'])
        self.listener = tf.TransformListener()

        self.state = TrafficLight.UNKNOWN
        self.last_state = TrafficLight.UNKNOWN
        self.last_wp = -1
        self.state_count = 0

        rospy.spin()

    def pose_cb(self, msg):
        self.pose = msg

    def waypoints_cb(self, waypoints):
        self.waypoints = waypoints
        # rospy.loginfo('Check is 2d waypoints created, else create it for KDTree')
        if not self.waypoints_2d:
            self.waypoints_2d = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in waypoints.waypoints]
            # rospy.loginfo('waypoints[0].pose.pose.position.x = %s', waypoints.waypoints[0].pose.pose.position.x)
            self.waypoint_tree = KDTree(self.waypoints_2d)
            # rospy.loginfo('waypoint_tree data = %s', self.waypoint_tree.data)        

    def traffic_cb(self, msg):
        self.lights = msg.lights

    def image_cb(self, msg):
        """Identifies red lights in the incoming camera image and publishes the index
            of the waypoint closest to the red light's stop line to /traffic_waypoint

        Args:
            msg (Image): image from car-mounted camera

        """
        self.has_image = True
        # Get time elapsed, then check has 200 milliseconds passed, so we can process image
        self.time_elapsed = default_timer() - self.last_image_processed
        if self.time_elapsed < CAMERA_IMAGE_PROCESSING_RATE:
            # rospy.loginfo("time_elapsed < CAMERA_IMAGE_PROCESSING_RATE: t.e. = %s", self.time_elapsed)
            return

        self.camera_image = msg
        self.last_image_processed = default_timer()
        light_wp, state = self.process_traffic_lights()
        # rospy.loginfo("tl_detector p1: light_wp = %s", light_wp)

        '''
        Publish upcoming red lights at camera frequency.
        Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
        of times till we start using it. Otherwise the previous stable state is
        used.
        '''
        if self.state != state:
            self.state_count = 0
            self.state = state
        elif self.state_count >= STATE_COUNT_THRESHOLD:
            self.last_state = self.state
            light_wp = light_wp if state == TrafficLight.RED else -1
            self.last_wp = light_wp
            # rospy.loginfo("tl_detector p2: self.light_wp = %s", self.last_wp)
            self.upcoming_red_light_pub.publish(Int32(light_wp))
        else:
            self.upcoming_red_light_pub.publish(Int32(self.last_wp))
        self.state_count += 1

    def get_closest_waypoint(self, x, y):
        """Identifies the closest path waypoint to the given position
            https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
        Args:
            pose (Pose): position to match a waypoint to

        Returns:
            int: index of the closest waypoint in self.waypoints

        """
        #TODO implement
        closest_idx = self.waypoint_tree.query([x, y], 1)[1]

        return closest_idx

    def get_light_state(self, light):
        """Determines the current color of the traffic light

        Args:
            light (TrafficLight): light to classify

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        # For testing if you don't have a detection model, just return the light state
        # return light.state

        # Using SSD MobileNet to Detect and Classify Traffic Lights

        if(not self.has_image):
            self.prev_light_loc = None
            return False

        cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")

        # Get classification
        save_classified_img = False
        return self.light_classifier.get_classification(cv_image, save_classified_img)

    def process_traffic_lights(self):
        """Finds closest visible traffic light, if one exists, and determines its
            location and color

        Returns:
            int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        closest_light = None
        line_wp_idx = None

        # List of positions that correspond to the line to stop in front of for a given intersection
        stop_line_positions = self.config['stop_line_positions']
        if(self.pose):
            car_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x, self.pose.pose.position.y)
            # rospy.loginfo("car_wp_idx = %s", car_wp_idx)

        #TODO find the closest visible traffic light (if one exists)
        diff = len(self.waypoints.waypoints)
        # rospy.loginfo("diff = %s", diff)
        for i, light in enumerate(self.lights):
            # Get stop line waypoint index
            line = stop_line_positions[i]
            temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
            # rospy.loginfo("temp_wp_idx = %s", temp_wp_idx)
            # Find closest stop line waypoint index
            d = temp_wp_idx - car_wp_idx
            # distance between closest waypoint in front of car and stop line
            # rospy.loginfo("distance d between car and stop line = %s", d)
            if d >= 0 and d < diff:
                diff = d
                closest_light = light
                line_wp_idx = temp_wp_idx
                # rospy.loginfo("line_wp_idx = %s", line_wp_idx)

        if closest_light and ( (line_wp_idx - car_wp_idx) <=  WAYPOINT_THRESHOLD):
            state = self.get_light_state(closest_light)
            # rospy.loginfo("return line_wp_idx = %s", line_wp_idx)
            return line_wp_idx, state

        # if the traffic light state is unknown: 
        # For this project, keep the car moving
        # udacity will be testing the real car in a controlled environment and if it
        # rolls through the traffic light, they can still test the drive-by-wire
        # system and make sure it is following waypoints
        # In real life, you would probably want to bring the car to a stop if it sees
        # a traffic light it can't determine the state of the traffic light
        return -1, TrafficLight.UNKNOWN

if __name__ == '__main__':
    try:
        TLDetector()
    except rospy.ROSInterruptException:
        rospy.logerr('Could not start traffic node.')