#!/usr/bin/env python import fileinput import random from kdtree import KDTree # read in the points from a file specified on the command line. E.g.: # $ ./kdtree_test.py ../../../data/sim_waypoints.csv points = [] i=0 for line in fileinput.input(): line_parts = line.split(',') points.append((float(line_parts[0]), float(line_parts[1]), int(i))) i += 1 # generate the K-D Tree from the points kdtree = KDTree(points) # pick a random point from the points rand_index = random.randint(0, len(points)) # find the closest point point = points[rand_index] print ("randomly chose point {} at index {}".format(point, rand_index)) new_point = (point[0] + 2.0, point[1] + 7.0) print ("tweaked x,y to be {}".format(new_point, rand_index)) closest = kdtree.closest_point(new_point) print ("closest point to {} is {}".format(new_point, closest))
class TLDetector(object): def __init__(self): rospy.init_node('tl_detector') self.pose = None self.waypoints = None self.camera_image = None 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) rospy.Subscriber('/image_color', Image, self.collect_images_callback) 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.listener = tf.TransformListener() self.state = TrafficLight.UNKNOWN self.last_state = TrafficLight.UNKNOWN self.last_wp = -1 self.state_count = 0 # Parameters for collecting frames from the camera self.should_collect_data = False self.dump_images_dir = create_dir_if_nonexistent( join(expanduser('~'), 'traffic_light_dataset', 'raw_images')) self.dump_images_counter = len(os.listdir(self.dump_images_dir)) self.last_dump_tstamp = rospy.get_time() # Used to find the closest waypoint self.kdtree = None # Data file to store the image name and light state in the image. self.datafile = open(self.dump_images_dir + "/lightsData.csv", "w+") self.lightState = None if not PREFER_GROUND_TRUTH: # Create tensorflow session self.session = tensorflow.Session() # Import classifier and restore pre-trained weights self.light_classifier = TrafficLightClassifier( input_shape=[64, 64], learning_rate=1e-4) tensorflow.train.Saver().restore( self.session, TrafficLightClassifier.checkpoint_path) rospy.spin() def collect_images_callback(self, msg): """ Save camera images (currently once per second) """ def should_collect_camera_image(): return self.should_collect_data and ( rospy.get_time() - self.last_dump_tstamp > 1) if should_collect_camera_image(): # Convert image message to actual numpy data image_data = self.bridge.imgmsg_to_cv2(msg) image_data = cv2.cvtColor( image_data, cv2.COLOR_RGB2BGR) # opencv uses BGR convention image_path = join(self.dump_images_dir, '{:06d}.jpg'.format(self.dump_images_counter)) # Dump image to dump directory cv2.imwrite(image_path, image_data) # write the state of the light and the image name to a csv file print("Writing to datafile") self.datafile.write('{:06d}.jpg'.format(self.dump_images_counter) + " , " + self.lightState + "\n") # Update counter and timestamp self.dump_images_counter += 1 self.last_dump_tstamp = rospy.get_time() 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 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 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 not None and self.kdtree is None: if VERBOSE: print('tl_detector: g_cl_wp: initializing kdtree') points = [] for i, waypoint in enumerate(self.waypoints): points.append((float(waypoint.pose.pose.position.x), float(waypoint.pose.pose.position.y), i)) self.kdtree = KDTree(points) if self.kdtree is not None: current_position = (pose.position.x, pose.position.y) closest = self.kdtree.closest_point(current_position) if VERBOSE: print('tl_detector: g_cl_wp: closest point to {} is {}'.format( current_position, closest)) return closest[2] return 0 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 not self.has_image: self.prev_light_loc = None return False cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8") light_state = self.light_classifier.get_classification( self.session, 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 closest 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 # 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_position = self.get_closest_waypoint(self.pose.pose) #TODO find the closest visible traffic light (if one exists) if VERBOSE: print("tl_detector: p_tl: There are {} traffic lights to analyze.". format(len(self.lights))) min_distance = float("Infinity") for current_light in self.lights: # Check to see whether the traffic light is ahead of the car if is_ahead(current_light, self.pose.pose): # Get the simplified Euclidean distance (no sqrt) between it and the car light_distance = get_simple_distance_from_waypoint( current_light, self.pose.pose) # If the light is closer, remember it if (light_distance < min_distance): min_distance = light_distance light = current_light # If we found a light ahead of us if light: self.lightState = self._light_color(light.state) # Calculate the actual distance the of the light. light_distance = math.sqrt(min_distance) if VERBOSE: print( "tl_detector: p_tl: closest light to {} is at {} (Distance: {})." .format( (self.pose.pose.position.x, self.pose.pose.position.y), (light.pose.pose.position.x, light.pose.pose.position.y), light_distance)) # Look up the closest waypoint to it # TODO: [brahm] Can we assume self.kdtree is initialized? light_wp = self.get_closest_waypoint(light.pose.pose) # Determine the state of the light state = -1 if PREFER_GROUND_TRUTH: if VERBOSE: print("tl_detector: p_tl: Ground truth light color: {}". format(self._light_color(light.state))) # TODO: [brahm] Determine what light.state is when not available (e.g. not in the simulator) if light.state is not None: state = light.state #if True: # (state == -1): if state == -1: # this is where we classify the light state_inferred = self.get_light_state(light) # If the traffic light is close, let us know if (light_distance < TL_NEARNESS_THRESHOLD): if VERBOSE: print("tl_detector: p_tl: light is close: {} meters away.". format(light_distance)) return light_wp, state self.waypoints = None return -1, TrafficLight.UNKNOWN # Helper def _light_color(self, state): if (state == TrafficLight.RED): return "RED" elif (state == TrafficLight.YELLOW): return "YELLOW" elif (state == TrafficLight.GREEN): return "GREEN" else: return "UNKNOWN"