Пример #1
0
class TLDetector(object):
    def __init__(self):
        rospy.init_node('tl_detector')

        self.pose = None
        self.waypoints = None
        self.camera_image = None
        self.lights = []
        self.waypoints_2D = None
        self.waypoint_TREE = None
        self.sim_testing = False
        self.gamma_correction = False
        self.count = 0
        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', GlobalTrafficLightMessage, queue_size=1)

        self.bridge = CvBridge()
        self.light_classifier = TLClassifier(0.01, 0.5, False)
        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

        if not self.waypoints_2D:
            self.waypoints_2D = [[
                waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
            ] for waypoint in waypoints.waypoints]
            self.waypoint_TREE = KDTree(self.waypoints_2D)

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

    def _pass_threshold(self):
        if self.state == TrafficLight.YELLOW:
            return self.state_count >= STATE_COUNT_THRESHOLD - 1
        return self.state_count >= STATE_COUNT_THRESHOLD

    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
        self.camera_image = msg

        light_wp, state = self.process_traffic_lights()
        '''
        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
        #     self.upcoming_red_light_pub.publish(Int32(light_wp))
        # else:
        #     self.upcoming_red_light_pub.publish(Int32(self.last_wp))
        # self.state_count += 1
        if self.state != state:
            self.state_count = 0
            self.state = state
        elif self._pass_threshold():
            self.last_wp = light_wp
            msg = self._prepare_result_msg(state, light_wp)
            #rospy.loginfo(msg)
            self.upcoming_red_light_pub.publish(msg)
        else:
            msg = self._prepare_result_msg(self.state, self.last_wp)
            rospy.loginfo(msg)
            self.upcoming_red_light_pub.publish(msg)
        self.state_count += 1

    def _prepare_result_msg(self, tl_state, tl_stop_waypoint):
        tl_result = GlobalTrafficLightMessage()
        tl_result.state = tl_state
        tl_result.waypoint = tl_stop_waypoint

        return tl_result

    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)

        """

        if self.light_classifier is None:  # Stop if classifer is not yet ready
            return TrafficLight.RED
        if (not self.has_image):
            self.prev_light_loc = None
            return False

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

        processed_img = cv_image[0:600, 0:800]  # was [20:400, 0:800]

        #Convert image to RGB format
        processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB)

        #Get classification

        #initialize light_state to unknown by default
        light_state = TrafficLight.UNKNOWN
        light_state_via_msg = None

        img_full_np = self.light_classifier.load_image_into_numpy_array(
            processed_img)

        unknown = False

        b, conf, cls_idx = self.light_classifier.get_localization_classification(
            img_full_np)

        if np.array_equal(b, np.zeros(4)):
            print('unknown')
            unknown = True
        else:
            #light_state = cls_idx
            if cls_idx == 1.0:
                print('Green', b)
                light_state = TrafficLight.GREEN
            elif cls_idx == 2.0:
                print('Red', b)
                light_state = TrafficLight.RED
            elif cls_idx == 3.0:
                print('Yellow', b)
                light_state = TrafficLight.YELLOW
            elif cls_idx == 4.0:
                print('Unknown', b)
                light_state = TrafficLight.UNKNOWN
            else:
                print('Really Unknown! Didn\'t process image well', b)
                light_state = TrafficLight.UNKNOWN

        self.count = self.count + 1
        #filename = "site_image_" + str(self.count) + str(light_state) + ".jpg"
        #cv2.imwrite(filename, cv_image)

        return light_state

    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)
            #TODO find the closest visible traffic light (if one exists)
            diff = len(self.waypoints.waypoints)

            isclosestTFavilable = false
            for i, light in enumerate(self.lights):
                line = stop_line_positions[i]
                temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
                d = temp_wp_idx - car_wp_idx

                if d >= 0 and d < diff:
                    rospy.loginfo(d)
                    diff = d
                    isclosestTFavilable = True
                    line_wp_idx = temp_wp_idx
                    if d < 67:
                        closest_light = light

        if closest_light and closest_light is not None:
            rospy.loginfo("Inside classifer")
            state = self.get_light_state(closest_light)
            return line_wp_idx, state

        if isclosestTFavilable:
            return line_wp_idx, TrafficLight.UNKNOWN
        else:
            return -1, TrafficLight.UNKNOWN
Пример #2
0
class TLDetector(object):
    def __init__(self):
        rospy.init_node('tl_detector')

        self.sim_testing = bool(rospy.get_param("~sim_testing", True))
        threshold = rospy.get_param('~threshold', 0.3)
        hw_ratio = rospy.get_param('~hw_ratio', 1.5)
        self.gamma_correction = bool(
            rospy.get_param("~gamma_correction", False))

        self.pose = None
        self.waypoints = None
        self.camera_image = None
        self.lights = []

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

        self.bridge = CvBridge()
        self.light_classifier = None
        self.light_classifier = TLClassifier(threshold, hw_ratio,
                                             self.sim_testing)
        self.listener = tf.TransformListener()

        img_full_np = self.light_classifier.load_image_into_numpy_array(
            np.zeros((800, 600, 3)))
        #prime the pump. Don't need to do it for get_localization_classification.

        self.light_classifier.get_localization(img_full_np)  # "prime the pump"

        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,
                                queue_size=1,
                                buff_size=2 * 52428800)  # extended

        #we don't use this publisher anymore. We publish the status of all traffic lights. Not just the red ones.
        self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
                                                      Int32,
                                                      queue_size=1)
        self.upcoming_traffic_light_pub = rospy.Publisher(
            '/all_traffic_waypoint', TLStatus, queue_size=1)

        self.state = None
        self.last_state = None
        self.last_wp = -1
        self.state_count = 0
        self.count = 0  #processed images count
        self.total_classification = 0
        self.tp_classification = 0.0
        print("TL Detection ... initialization complete")

        rospy.spin()

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

    def waypoints_cb(self, waypoints):
        self.waypoints = waypoints.waypoints

    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

        # rospy.loginfo('Image seq %s', msg.header.seq)
        # img_time = msg.header.stamp
        # age = rospy.Time.now() - img_time
        # if age.secs > 0.2:
        #     rospy.loginfo('Image too old. Age %s', age.secs)
        #     light_wp, state = -1, TrafficLight.UNKNOWN
        # else:
        self.camera_image = msg
        light_wp, state = self.process_traffic_lights()
        '''
        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 = state
            #light_wp = light_wp if state == TrafficLight.RED else -1
            self.last_wp = light_wp
            tl_status_msg = TLStatus()
            tl_status_msg.header.frame_id = '/world'
            tl_status_msg.header.stamp = rospy.Time.now()
            tl_status_msg.waypoint = light_wp
            tl_status_msg.state = state
            # rospy.loginfo("Light Wp is %s and state is %s", light_wp, self.state)
            self.upcoming_traffic_light_pub.publish(tl_status_msg)
        else:
            # we keep publishing the last state and wp until it gets confirmed.
            tl_status_msg = TLStatus()
            tl_status_msg.header.frame_id = '/world'
            tl_status_msg.header.stamp = rospy.Time.now()
            tl_status_msg.waypoint = self.last_wp
            if self.last_state is not None:
                tl_status_msg.state = self.last_state
            else:
                # Do not start moving until first light state is obtained
                # Fix for starting issues
                tl_status_msg.state = TrafficLight.RED
            #rospy.loginfo("Light Wp is %s and state is %s:", self.last_wp, self.last_state)
            self.upcoming_traffic_light_pub.publish(tl_status_msg)
        self.state_count += 1

    def get_closest_waypoint(self, pose):
        """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

        """
        if self.waypoints is None:
            return
        #TODO implement
        min_dist = 10000
        min_loc = None

        pos_x = pose.position.x
        pos_y = pose.position.y
        # check all the waypoints to see which one is the closest to our current position
        for i, waypoint in enumerate(self.waypoints):
            wp_x = waypoint.pose.pose.position.x
            wp_y = waypoint.pose.pose.position.y
            dist = math.sqrt((pos_x - wp_x)**2 + (pos_y - wp_y)**2)
            if (dist < min_dist):  #we found a closer wp
                min_loc = i  # we store the index of the closest waypoint
                min_dist = dist  # we save the distance of the closest waypoint

        # returns the index of the closest waypoint
        return min_loc

    def find_stop_line(self, closest_light_wp):
        """Find waypoint for stop line associated to traffic light
        Args:
            closest_light_wp: closest waypoint to the light
        Returns:
            closest_stop_line_wp: waypoint closest to the sop line before the light
        """
        stop_line_positions = self.config['stop_line_positions']
        closest_light_stop_wp = None

        #search  the stop line waypoint which is the closest to closest_light_wp
        min_dist = 10000
        for light_stop_position in stop_line_positions:
            #convert 2D position into Pose format to get closest waypoint
            light_stop_pose = Pose()
            light_stop_pose.position.x = light_stop_position[0]
            light_stop_pose.position.y = light_stop_position[1]
            light_stop_wp = self.get_closest_waypoint(light_stop_pose)
            dist = abs(closest_light_wp - light_stop_wp)
            if dist < min_dist:
                min_dist = dist
                closest_light_stop_wp = light_stop_wp  #note that it can be a bit ahead of the stop line

        return closest_light_stop_wp

    def project_to_image_plane(self, point_in_world):
        """Project point from 3D world coordinates to 2D camera image location
            Args:
                point_in_world (Point): 3D location of a point in the world
                Returns:
                    x (int): x coordinate of target point in image
                    y (int): y coordinate of target point in image
        """
        fx = self.config['camera_info']['focal_length_x']
        fy = self.config['camera_info']['focal_length_y']
        image_width = self.config['camera_info']['image_width']
        image_height = self.config['camera_info']['image_height']

        #print("Image Size:", image_width, image_height)
        cx = image_width / 2
        cy = image_height / 2

        # get transform between pose of camera and world frame
        trans = None
        rot = None
        x = 0
        y = 0
        try:
            now = rospy.Time.now()
            self.listener.waitForTransform("/base_link", "/world", now,
                                           rospy.Duration(1.0))
            (trans,
             rot) = self.listener.lookupTransform("/base_link", "/world", now)

        except (tf.Exception, tf.LookupException, tf.ConnectivityException):
            rospy.logerr("Failed to find camera to map transform")
            return None, None

        if (trans and rot):
            rpy = tf.transformations.euler_from_quaternion(rot)
            yaw = rpy[2]

            (ptx, pty,
             ptz) = (point_in_world.position.x, point_in_world.position.y,
                     point_in_world.position.z)

            #rotation
            point_to_cam = (ptx * math.cos(yaw) - pty * math.sin(yaw),
                            ptx * math.sin(yaw) + pty * math.cos(yaw), ptz)
            #translation
            point_to_cam = [sum(x) for x in zip(point_to_cam, trans)]

            #project to plan
            x = -fx * point_to_cam[1] / point_to_cam[0]
            y = -fy * point_to_cam[2] / point_to_cam[0]

            x = int(x + cx)
            y = int(y + cy)

        return (x, y)

    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)

        """

        if self.light_classifier is None:  # Stop if classifer is not yet ready
            return TrafficLight.RED
        if (not self.has_image):
            self.prev_light_loc = None
            return False

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

        #TODO use light location to zoom in on traffic light in image
        #Prepare image for classification
        if self.sim_testing:  #we cut 50 pixels left and right of the image and the bottom 100 pixels
            width, height, _ = cv_image.shape
            x_start = int(width * 0.10)
            x_end = int(width * 0.90)
            y_start = 0
            y_end = int(height * 0.85)
            processed_img = cv_image[y_start:y_end, x_start:x_end]
        else:  # real-case testing. Reduce image size to avoid light reflections on hood.
            processed_img = cv_image[0:600, 0:800]  # was [20:400, 0:800]

        #Convert image to RGB format
        processed_img = cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB)

        #Get classification

        #initialize light_state to unknown by default
        light_state = TrafficLight.UNKNOWN
        light_state_via_msg = None

        #get the ground truth traffic light states through the traffic light messages
        for tl in self.lights:
            dist = math.sqrt((tl.pose.pose.position.x - light.position.x)**2 +
                             (tl.pose.pose.position.y - light.position.y)**2)
            if (dist < 50):  #means we found the light close to the stop line
                light_state_via_msg = tl.state
                break  #no need to parse other lights once light was found

        #detect traffic light position (box) in image
        #convert image to np array
        img_full_np = self.light_classifier.load_image_into_numpy_array(
            processed_img)

        #apply gamma correction to site testing if parameter set in launch file.
        if (self.gamma_correction == True):
            img_full_np = adjust_gamma(img_full_np, 0.4)

        # if simulator, we apply detection and classification separately

        unknown = False

        if self.sim_testing:
            # find traffic light in image.
            b = self.light_classifier.get_localization(img_full_np)
            print(b)
            # If there is no detection or low-confidence detection
            if np.array_equal(b, np.zeros(4)):
                print('unknown')
                unknown = True
            else:  #we can use the classifier to classify the state of the traffic light
                img_np = cv2.resize(processed_img[b[0]:b[2], b[1]:b[3]],
                                    (32, 32))
                self.light_classifier.get_classification(img_np)
                light_state = self.light_classifier.signal_status
        else:
            print("Get in Localization-Classification")
            b, conf, cls_idx = self.light_classifier.get_localization_classification(
                img_full_np, visual=False)
            print("Get out of Localization-Classification")
            if np.array_equal(b, np.zeros(4)):
                print('unknown')
                unknown = True
            else:
                #light_state = cls_idx
                if cls_idx == 1.0:
                    print('Green', b)
                    light_state = TrafficLight.GREEN
                elif cls_idx == 2.0:
                    print('Red', b)
                    light_state = TrafficLight.RED
                elif cls_idx == 3.0:
                    print('Yellow', b)
                    light_state = TrafficLight.YELLOW
                elif cls_idx == 4.0:
                    print('Unknown', b)
                    light_state = TrafficLight.UNKNOWN
                else:
                    print('Really Unknown! Didn\'t process image well', b)
                    light_state = TrafficLight.UNKNOWN

        #check prediction against ground truth
        if self.sim_testing:
            rospy.loginfo("Upcoming light %s, True state: %s", light_state,
                          light_state_via_msg)
            #compare detected state against ground truth for (simulator only)
            if not unknown:
                self.count = self.count + 1
                filename = "sim_image_" + str(self.count)
                if (light_state == light_state_via_msg):
                    self.tp_classification = self.tp_classification + 1
                    #filename = filename + "_good_" + str(light_state) + ".jpg"
                else:
                    filename = filename + "_bad_" + str(light_state) + ".jpg"
                    #cv2.imwrite(filename, cv_image)
                self.total_classification = self.total_classification + 1
                accuracy = (self.tp_classification /
                            self.total_classification) * 100
                if self.count % 20 == 0:
                    rospy.loginfo("Classification accuracy: %s", accuracy)
        else:  #site testing
            self.count = self.count + 1
            #filename = "site_image_" + str(self.count) + str(light_state) + ".jpg"
            #cv2.imwrite(filename, cv_image)

        return light_state

    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)

        """
        light = None
        closest_light_wp = None
        closest_light_stop_wp = None
        dist_to_light = 10000  #initialize to high value

        stop_line_positions = self.config['stop_line_positions']
        if (self.pose):
            car_position = self.get_closest_waypoint(self.pose.pose)
            #rospy.loginfo("Car position (at Wp Index): %s", car_position)
        else:
            return -1, TrafficLight.UNKNOWN

        for light_stop_position in stop_line_positions:
            light_stop_pose = Pose()
            light_stop_pose.position.x = light_stop_position[0]
            light_stop_pose.position.y = light_stop_position[1]
            light_stop_wp = self.get_closest_waypoint(
                light_stop_pose)  #get the wp closest to each light_position
            if light_stop_wp >= car_position:  #it found a waypoint close to the traffic light and ahead of the car
                if closest_light_stop_wp is None:  #check if this is the first light we process
                    closest_light_stop_wp = light_stop_wp
                    light = light_stop_pose
                elif light_stop_wp < closest_light_stop_wp:
                    closest_light_stop_wp = light_stop_wp  #if we have a closer light_wp ahead of the car position, we allocate closer value
                    light = light_stop_pose

        if ((car_position is not None)
                and (closest_light_stop_wp is not None)):
            dist_to_light = abs(car_position - closest_light_stop_wp)
            #rospy.loginfo("Closest light position (in Wp index): %s", closest_light_stop_wp)

        if light and dist_to_light < 100:  #we check the status of the traffic light if it's within 200 waypoints distance
            state = self.get_light_state(light)
            #closest_stop_line_wp = self.find_stop_line(closest_light_wp)
            return closest_light_stop_wp, state
        #self.waypoints = None
        return -1, TrafficLight.UNKNOWN