class TLDetector(object): def __init__(self): # ============== ROS specific stuff rospy.init_node('tl_detector') rospy.Subscriber('/current_pose', PoseStamped, self.cb_pose) rospy.Subscriber('/base_waypoints', Lane, self.cb_waypoints) rospy.Subscriber('/idx_closest_waypoint', std_msgs.msg.Int32, self.cb_waypoint_next) rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray, self.cb_traffic_lights) rospy.Subscriber('/image_color', Image, self.cb_image) self.pub_idx_wp_to_stop = rospy.Publisher('/traffic_waypoint', Int32, queue_size=1) # ========== other stuff # waypoints and egopose self.pose = None self.waypoints_2d = None self.waypoints_tree = None # traffic lights / stop lines config_string = rospy.get_param("/traffic_light_config") self.config = yaml.load(config_string) self.lights = None self.idx_light_next = 0 # image classification self.cv_bridge = CvBridge() self.img_cnt = 0 self.img_t_last = time.time() self.classifier = TLClassifier(filename_pb="model.pb") self.img_queue = [] # start looping rospy.spin() def cb_pose(self, msg): self.pose = msg def cb_waypoints(self, msg_waypoints): if self.waypoints_2d is None: # only do it once! # collect waypoints self.waypoints_2d = [[ wp.pose.pose.position.x, wp.pose.pose.position.y ] for wp in msg_waypoints.waypoints] self.waypoints_tree = scipy.spatial.KDTree(self.waypoints_2d) # collect information about all traffic lights from config and waypoints self.lights = [] lights_pos_xy = self.config['stop_line_positions'] for idx_light, pos_xy in enumerate(lights_pos_xy): light = Light() light.idx = idx_light light.x = pos_xy[0] light.y = pos_xy[1] light.idx_wp = self.waypoints_tree.query(pos_xy, 1)[1] self.lights.append(light) def cb_waypoint_next(self, msg): idx_wp_next = msg.data # check whether next light is still ahead of current waypoint. Otherwise, get next light if self.lights is not None: light_next = self.lights[self.idx_light_next] if idx_wp_next > light_next.idx_wp: self.idx_light_next = (self.idx_light_next + 1) % len( self.lights) def cb_traffic_lights(self, msg): if self.lights is not None: # copy information from msg into class object for idx, light_gt in enumerate(msg.lights): self.lights[idx].state_true = light_gt.state def cb_image( self, msg_img, num_images_skip=3, dt_min_between_images=0.050, flag_export=False, ): # only do stuff with image if next traffic light is close. Also don't process every image distance = self.get_distance_to_next_light() if distance is not None and distance < 100: self.img_cnt += 1 time_now = time.time() if self.img_cnt % num_images_skip == 0 and time_now - self.img_t_last > dt_min_between_images: self.img_t_last = time_now # start of actual processing img_numpy = self.cv_bridge.imgmsg_to_cv2(msg_img, "bgr8") if flag_export: self.export_image(img_numpy, distance) else: state = self.predict_state(img_numpy) if state is not None: next_light = self.lights[self.idx_light_next] next_light.state_pred = state self.check_lights() def predict_state( self, img, batch_size=3, target_size=(320, 240), ratio_min_detection=0.6, ): # resize image, scale values, convert to RGB and and append to queue img = cv2.resize(img, target_size, interpolation=cv2.INTER_LINEAR) img = img / 255. img = img[:, :, ::-1] # converts BGR to RGB self.img_queue.append(img) # only if queue has batch_size items in it, run prediction if len(self.img_queue) == batch_size: input_tensor = np.asarray(self.img_queue) assert (input_tensor.ndim == 4) self.img_queue = [] prob_all = self.classifier.predict(input_tensor) prob_red = prob_all[:, 0] if np.mean(prob_red > 0.5) > ratio_min_detection: return 0 # meaning traffic light = red elif np.mean(prob_red < 0.5) > ratio_min_detection: return 2 # meaning traffic light = green else: return None else: return None def export_image(self, img_numpy, distance): light_state = self.lights[ self.idx_light_next].state_true # state as int 0,1,2,3, see above time_in_ms = np.round(time.time() * 1000).astype(np.int) folder = os.path.join('/mnt/share/export', str(light_state)) filename = ( 'img_' + str(time_in_ms) # + '_state_' + str(light_state) + '_lightidx_' + str(self.idx_light_next) + '_dist_' + str(np.round(distance).astype(np.int)) + '.png') filepath = os.path.join(folder, filename) cv2.imwrite(filepath, img_numpy) rospy.loginfo("Written image {}".format(filename)) def check_lights(self): msg_out = Int32() msg_out.data = -1 # default value, meaning no need to stop at any waypoint distance = self.get_distance_to_next_light() if distance is not None: light_next = self.lights[self.idx_light_next] if light_next.state_pred == 0: # 0=red, see above msg_out.data = light_next.idx_wp if True: rospy.loginfo( "Light idx={}: state_pred={}, state_true={}, distance={}". format( self.idx_light_next, COLOR_PER_INT[light_next.state_pred], COLOR_PER_INT[light_next.state_true], distance, )) self.pub_idx_wp_to_stop.publish(msg_out) def get_distance_to_next_light(self): if self.pose and self.lights and self.idx_light_next is not None: car_xy = np.asarray( [self.pose.pose.position.x, self.pose.pose.position.y]) light_next = self.lights[self.idx_light_next] light_next_xy = np.asarray([light_next.x, light_next.y]) distance = self.calc_distance(car_xy, light_next_xy) return distance else: return None @classmethod def calc_distance(self, pt1, pt2): dist = pt1 - pt2 dist_norm = np.linalg.norm(dist) return dist_norm
class TLDetector(object): """ detect and classify traffic light @subscribed /base_waypoints: the complete list of waypoints the car will be following @subscribed /current_pose: the vehicle's current position @subscribed /image_color: the image stream from the car's camera @subscribed /vehicle/traffic_lights: the exact location and status of all traffic lights in simulator @published /traffic_waypoint: the index of the waypoint for nearest upcoming red light's stop line """ CAMERA_IMAGE_CLASSIFICATION_WPS = 53 CAMERA_IMAGE_COLLECTION_AFTER_LINE_COUNT = 27 def __init__(self): rospy.init_node('tl_detector') # load config params: config_string = rospy.get_param("/traffic_light_config") self.config = yaml.load(config_string) # state variables: self.pose = None self.waypoints = None self._waypoints_location = None self._waypoints_size = None self._waypoints_index = None self.camera_image = None self.lights = [] self.state = TrafficLight.UNKNOWN self.last_state = TrafficLight.UNKNOWN self.last_wp = -1 self.state_count = 0 # image collector: self.after_stop_line_count = TLDetector.CAMERA_IMAGE_COLLECTION_AFTER_LINE_COUNT # classifier--subscriber: self.listener = tf.TransformListener() # classifier--format convertor: self.bridge = CvBridge() # classifier--pre-trained model: filenames = os.listdir('./light_classification/models') if not filenames: pass else: # model name pattern: FILENAME_PATTERN = re.compile('(\d+)-model-params.h5') # parse model timestamps: timestamps = [ int(FILENAME_PATTERN.match(filename).group(1)) for filename in filenames ] # identify latest model: _, latest_model_filename = max(zip(timestamps, filenames), key=lambda t: t[0]) # load latest model: self.light_classifier = TLClassifier() self.light_classifier.load( os.path.join('./light_classification/models', latest_model_filename)) # subscribe: rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb) rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb) rospy.Subscriber('/image_color', Image, self.image_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. ''' rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray, self.traffic_cb) # publish: self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint', Int32, queue_size=1) rospy.spin() def waypoints_cb(self, waypoints): """ load base waypoints from system Args: waypoints (list of Waypoint): reference trajectory as waypoints """ if not self.waypoints: # load waypoints: self.waypoints = waypoints # build index upon waypoints: self._waypoints_location = np.array([[ waypoint.pose.pose.position.x, waypoint.pose.pose.position.y ] for waypoint in waypoints.waypoints]) self._waypoints_size, _ = self._waypoints_location.shape self._waypoints_index = KDTree(self._waypoints_location) def pose_cb(self, msg): """ parse ego vehicle pose Args: msg (PoseStamped): ego vehicle pose """ self.pose = msg def image_cb(self, msg): """ identify 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 """ # parse input: self.has_image = True self.camera_image = msg # process traffic lights: stop_line_waypoint_index, 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 stop_line_waypoint_index = stop_line_waypoint_index if ( state == TrafficLight.RED or state == TrafficLight.YELLOW) else -1 self.last_wp = stop_line_waypoint_index self.upcoming_red_light_pub.publish( Int32(stop_line_waypoint_index)) else: self.upcoming_red_light_pub.publish(Int32(self.last_wp)) self.state_count += 1 def traffic_cb(self, msg): """ parse traffic light status from telegram Args: msg (TrafficLightArray): list of all traffic light telegrams """ self.lights = msg.lights def get_next_waypoint_index(self): """ get next waypoint index for ego vehicle """ # ego vehicle location: ego_vehicle_location = np.array( [self.pose.pose.position.x, self.pose.pose.position.y]) # closest waypoint _, closest_waypoint_index = self._waypoints_index.query( ego_vehicle_location) closest_waypoint_location = self._waypoints_location[ closest_waypoint_index] previous_waypoint_location = self._waypoints_location[ closest_waypoint_index - 1] # whether the closest waypoint is the next waypoint: is_next_waypoint = (np.dot( closest_waypoint_location - ego_vehicle_location, previous_waypoint_location - ego_vehicle_location) < 0.0) # next waypoint index: next_waypoint_index = closest_waypoint_index if not is_next_waypoint: next_waypoint_index = (closest_waypoint_index + 1) % self._waypoints_size return next_waypoint_index def get_closest_waypoint(self, position): """ get closest waypoint index for stop line Args: position (Pose): ego vehicle pose """ location = np.array(position) _, index = self._waypoints_index.query(location) return index def get_light_state_from_camera(self): """ Determines the closest traffic light state from image analysis """ if (not self.has_image): self.prev_light_loc = None return TrafficLight.UNKNOWN # format as OpenCV: cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8") # preprocess: preprocessed_image = self.light_classifier.preprocess(cv_image) # predict: return self.light_classifier.predict(preprocessed_image[np.newaxis]) def save_traffic_light_image(self, index, order, distance, state): """ Save traffic light image for offline training Args: index (Int): traffic light index order (str): 'before' or 'after' distance (Int): distance to incoming stop line state (TrafficLight.state): traffic light state """ # format image: traffic_light_image = self.bridge.imgmsg_to_cv2( self.camera_image, "bgr8") preprocessed = self.light_classifier.preprocess(traffic_light_image) filename = "light_classification/traffic_light_images/{}--{}-{}--{}=={}-preprocessed.jpg".format( rospy.Time.now().to_nsec(), order, index, distance, state) cv2.imwrite(filename, preprocessed) def get_light_state_from_telegram(self, light): """ Determines the closest traffic light state from telegram broadcast Args: light (TrafficLight): traffic light status """ 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_distance = self._waypoints_size closest_stop_line_index = None closest_stop_line_waypoint_index = 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) and self.waypoints: # ego vehicle position: ego_vehicle_waypoint_index = self.get_next_waypoint_index() # identify closest stop line: for i, stop_line_position in enumerate(stop_line_positions): stop_line_waypoint_index = self.get_closest_waypoint( stop_line_position) distance = stop_line_waypoint_index - ego_vehicle_waypoint_index if distance > 0 and distance < closest_distance: closest_distance = distance closest_stop_line_index = i closest_stop_line_waypoint_index = stop_line_waypoint_index # if there is incoming stop line: if ((closest_stop_line_waypoint_index) and (closest_distance <= TLDetector.CAMERA_IMAGE_CLASSIFICATION_WPS or self.after_stop_line_count > 0)): order = "before" # ego vehicle just passed stop line: if closest_distance > TLDetector.CAMERA_IMAGE_CLASSIFICATION_WPS: if self.after_stop_line_count > 0: self.after_stop_line_count -= 1 order = "after" closest_stop_line_index -= 1 closest_distance = self.after_stop_line_count # ego vehicle is about to cross stop line: elif closest_distance <= 3 and self.after_stop_line_count <= 0: self.after_stop_line_count = TLDetector.CAMERA_IMAGE_COLLECTION_AFTER_LINE_COUNT # method 01: telegram: state_telegram = self.get_light_state_from_telegram( self.lights[closest_stop_line_index]) # method 02: image analysis state_camera = self.get_light_state_from_camera() # image collection: if state_telegram != state_camera and state_camera != TrafficLight.UNKNOWN: # save for hard negative mining: self.save_traffic_light_image(closest_stop_line_index, order, closest_distance, state_telegram) # prompt: rospy.logwarn( "[Discrepancy between Telegram and Camera]: %d--%d @ %d, Camera Image Saved", state_telegram, state_camera, closest_stop_line_index) return closest_stop_line_waypoint_index, state_camera ''' # save for hard negative mining: self.save_traffic_light_image( 0, "unknown", 0, 0 ) rospy.logwarn("test site image saved.") ''' return -1, TrafficLight.UNKNOWN