class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.camera_image = None
self.lights = []
# auxiliary storage for waypoints, to allow for faster search of closest point
self.waypoints_2d = None
self.waypoint_tree = None
'''
/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.
'''
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.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.traffic_light_waypoints = None
self.has_image = False
rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
rospy.Subscriber('/image_color', Image, self.image_cb)
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 = cKDTree(self.waypoints_2d, leafsize=1)
self.traffic_light_waypoints = [
self.get_closest_waypoint(stop_line[0], stop_line[1], False)
for stop_line in self.config['stop_line_positions']
]
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 = -1
state = TrafficLight.UNKNOWN
if self.pose and self.waypoint_tree:
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 or state == TrafficLight.YELLOW 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, x, y, forward):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
x: x-coordinate of position to match a waypoint to
y: y-coordinate of position to match a waypoint to
forward: if true then find a closest waypoint ahead of (x,y)
if false then find a closest waypoint behind (x,y)
Returns:
int: index of the closest waypoint in self.waypoints
"""
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
# Check if closest is ahead or behind vehicle
closest_coord = self.waypoints_2d[closest_idx]
prev_coord = self.waypoints_2d[closest_idx - 1]
# Equation for hyperplane through closest_coords
cl_vect = np.array(closest_coord)
prev_vect = np.array(prev_coord)
pos_vect = np.array([x, y])
val = np.dot(cl_vect - prev_vect, pos_vect - cl_vect)
if val > 0 and forward:
closest_idx = (closest_idx + 1) % len(self.waypoints_2d)
elif val < 0 and not forward:
closest_idx = (closest_idx - 1) % len(self.waypoints_2d)
return closest_idx
def get_light_state(self):
"""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:
return TrafficLight.RED
if not self.has_image:
return TrafficLight.UNKNOWN
if hasattr(self.camera_image, 'encoding'):
if self.camera_image.encoding == '8UC3':
self.camera_image.encoding = "rgb8"
else:
self.camera_image.encoding = 'rgb8'
input_image = self.bridge.imgmsg_to_cv2(self.camera_image, "rgb8")
light_state = self.light_classifier.get_class(input_image)
#print("predicted state: ", light_state)
#print("ground thruth: ", light.state)
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)
"""
car_idx = self.get_closest_waypoint(self.pose.pose.position.x,
self.pose.pose.position.y, True)
# find the closest visible traffic light (if one exists)
diff = len(self.waypoints.waypoints)
light_idx = None
for temp_idx in self.traffic_light_waypoints:
# Find closest stopline waypoint index
dist = temp_idx - car_idx
if dist >= 0 and dist < diff:
diff = dist
light_idx = temp_idx
if light_idx and diff < 200:
state = self.get_light_state()
return light_idx, state
return -1, TrafficLight.UNKNOWN
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.has_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)
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.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 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
# print("*** RED LIGHT! *** ")
# print("State Count: " + str(self.state_count))
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, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
x: longitudinal position to match a waypoint to
y: lateral position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
closest_idx = -1
if self.waypoint_tree is not None:
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
# check if closest point is ahead or behind of ego
closest_coord = self.waypoints_2d[closest_idx]
prev_coord = self.waypoints_2d[closest_idx - 1]
# hyperplane eqn thru closest coords
cl_vect = np.array(closest_coord)
prev_vect = np.array(prev_coord)
pos_vect = np.array([x, y])
val = np.dot(cl_vect - prev_vect, pos_vect - cl_vect)
if val > 0:
closest_idx = closest_idx + 1 % len(self.waypoints_2d)
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:
return TrafficLight.RED
if not self.has_image:
return TrafficLight.UNKNOWN
if hasattr(self.camera_image, 'encoding'):
if self.camera_image.encoding == '8UC3':
self.camera_image.encoding = "rgb8"
else:
self.camera_image.encoding = 'rgb8'
img = self.bridge.imgmsg_to_cv2(self.camera_image, "rgb8")
light_state = self.light_classifier.get_class(img)
#print("predicted state: ", light_state)
#print("ground thruth: ", light.state)
return light_state #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 and self.waypoints:
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)
for i, light in enumerate(self.lights):
# get stop line wp idx
line = stop_line_positions[i]
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
# find closest stop line waypoint idx
d = temp_wp_idx - car_wp_idx
if d >= 0 and d < diff:
diff = d
closest_light = light
line_wp_idx = temp_wp_idx
if closest_light:
state = self.get_light_state(closest_light)
return line_wp_idx, state
return -1, TrafficLight.UNKNOWN