class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector', log_level=rospy.DEBUG)
self.pose = None
self.base_waypoints = None
self.waypoints_2d = None
self.waypoint_tree = None
self.camera_image = None
self.lights = []
self.number_of_detected_lights = 0
self.has_image = False
self.thread_working = False
self.frame_count = 0
self.bridge = CvBridge()
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.light_classifier = TLClassifier()
self.state = TrafficLight.UNKNOWN
self.debounced_state = TrafficLight.UNKNOWN
self.debounced_stop_wp_idx = -1
self.state_count = 0
self.stop_line_positions = self.config['stop_line_positions']
self.stop_line_wpidx = []
self.system_ready_flag = False
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane,
self.base_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_lights_cb)
sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
rospy.spin()
def pose_cb(self, msg):
self.pose = msg
def base_waypoints_cb(self, waypoints):
self.base_waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[
wp.pose.pose.position.x, wp.pose.pose.position.y
] for wp in waypoints.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
for i in range(len(self.stop_line_positions)):
line = self.stop_line_positions[i]
self.stop_line_wpidx.append(
self.get_closest_waypoint(line[0], line[1]))
rospy.logdebug("[tl_detector] base_waypoints received, length:%d",
len(self.base_waypoints.waypoints))
rospy.logdebug(
"[tl_detector] stop_line_waypoints processed, length:%d",
len(self.stop_line_wpidx))
def detect_tl(self):
#rospy.loginfo("Detection start")
stop_wp_idx, 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.
'''
car_wp_idx = -1
if (self.pose is not None) and (self.base_waypoints is not None):
car_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x,
self.pose.pose.position.y)
if self.state != state:
self.state_count = 0
self.state = state
else:
self.state_count += 1
if self.state_count >= STATE_COUNT_THRESHOLD:
if self.debounced_state != self.state:
rospy.logwarn(
"[tl_detector] Debounced light state change: %s -> %s ",
TRAFFIC_LIGHT_NAME[self.debounced_state],
TRAFFIC_LIGHT_NAME[self.state])
self.debounced_state = self.state
if self.state == TrafficLight.GREEN or state == 3:
self.debounced_stop_wp_idx = -1
else:
self.debounced_stop_wp_idx = stop_wp_idx
self.upcoming_red_light_pub.publish(Int32(self.debounced_stop_wp_idx))
self.thread_working = False
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
"""
#rospy.loginfo("Image_cb")
self.has_image = True
if not self.thread_working:
self.thread_working = True
self.camera_image = msg
thread.start_new_thread(self.detect_tl, ())
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
if self.waypoint_tree != None:
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
else:
closest_idx = -1
return closest_idx
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)
"""
stop_wp_idx = -1
state = TrafficLight.UNKNOWN
if (self.pose
is not None) and (self.has_image) and (self.base_waypoints
is not None):
car_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x,
self.pose.pose.position.y)
start = time.time()
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
self.number_of_detected_lights = self.light_classifier.detect_traffic_lights(
cv_image)
end1 = time.time()
if self.number_of_detected_lights > 0:
state = self.light_classifier.get_classification()
rospy.logdebug(
"[tl_detector] Detection Time:%.4fs, num of lights:%d, light_state:%s",
end1 - start, self.number_of_detected_lights,
TRAFFIC_LIGHT_NAME[state])
shortest_dist = len(self.base_waypoints.waypoints)
for i in range(len(self.stop_line_wpidx)):
d = self.stop_line_wpidx[i] - car_wp_idx
if d >= 0 and d < shortest_dist:
shortest_dist = d
stop_wp_idx = self.stop_line_wpidx[i]
else:
state = 3
rospy.logdebug("[tl_detector] No trafic light found!")
return stop_wp_idx, state
#!/usr/bin/env python
from light_classification.tl_classifier import TLClassifier
import os
import sys
import time
import cv2
if __name__ == '__main__':
light_classifier = TLClassifier()
while True:
filename = input("Filename:")
if filename == 'q':
sys.exit(0)
start = time.time()
img_filename = './00000013.png'
test_image = cv2.imread('./00000013.png')
output_img = light_classifier.detect_traffic_lights(test_image)
end = time.time()
print("Time: %s s" % (end - start))
