class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.camera_image = None
self.lights = []
# Waypoint KD Tree
self.waypoints_2d = None
self.waypoint_tree = None
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)
sub4 = rospy.Subscriber('/image_color', Image, self.image_cb)
# darknet_ros message
sub5 = rospy.Subscriber('/darknet_ros/bounding_boxes', BoundingBoxes,
self.detected_bb_cb)
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
# Get simulator_mode parameter (1== ON, 0==OFF)
self.simulator_mode = rospy.get_param("/simulator_mode")
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
# WARNING! Only use during testing in site mode (for diagnostics)
if int(self.simulator_mode) == 0:
self.cropped_tl_bb_pub = rospy.Publisher('/cropped_bb',
Image,
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
self.TL_BB_list = None
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
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 detected_bb_cb(self, msg):
# Clear the list
self.TL_BB_list = []
# Parameters: Diagnonal size thresholds
simulator_bb_size_threshold = 85 #px
site_bb_size_threshold = 40 #px
# min probability of detection
simulator_bb_probability = 0.85
site_bb_probability = 0.25
if int(self.simulator_mode) == 1:
prob_thresh = simulator_bb_probability
size_thresh = simulator_bb_size_threshold
else:
prob_thresh = site_bb_probability
size_thresh = site_bb_size_threshold
for bb in msg.bounding_boxes:
# Simulator mode: Bounding Box class should be 'traffic light' with probability >= 85%
# Site Mode: Bounding Box class should be 'traffic light' with probability >= 25%
if str(bb.Class
) == 'traffic light' and bb.probability >= prob_thresh:
# Simulator mode: If diagonal size of bounding box is more than 85px
# Site mode: If diagonal size of bounding box is more than 80px
if math.sqrt((bb.xmin - bb.xmax)**2 +
(bb.ymin - bb.ymax)**2) >= size_thresh:
self.TL_BB_list.append(bb)
# if running in site mode/ROS bag mode
if int(self.simulator_mode) == 0:
'''The ROS bag version only has video data. Hence no waypoints are loaded and get light function is not called.
So to check detection in ROS bag video, we do TL state classification here itself.
'''
# Get the camera image
cv_image = self.bridge.imgmsg_to_cv2(
self.camera_image, "bgr8")
# Crop image
bb_image = cv_image[bb.ymin:bb.ymax, bb.xmin:bb.xmax]
self.light_classifier.detect_light_state(bb_image)
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
"""
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 (not self.has_image):
self.prev_light_loc = None
return False
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
#Get classification
return self.light_classifier.get_classification(
cv_image, self.TL_BB_list, self.simulator_mode)
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)
diff = len(self.waypoints.waypoints)
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:
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
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.UseKnownTL = False
self.pose = None
self.waypoints = None
self.waypoints_2d = None
self.waypoint_tree = None
self.camera_image = None
#self.detecting_camera_image_list = []
#self.detecting_camera_image_list_orig = []
self.darknet_detecting_camera_image = None
self.detecting_camera_image = None
self.has_new_detecting_image = False
self.lights = []
self.images_path = "/home/louie/"
self.detected_img_count = 0
#self.darknet_img_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)
sub7 = rospy.Subscriber('/darknet_ros/bounding_boxes', BoundingBoxes,
self.darknet_bboxes_cb)
sub8 = rospy.Subscriber('/darknet_ros/detection_image', Image,
self.darknet_detected_image_cb)
sub8 = rospy.Subscriber('/darknet_ros/found_object', Int8,
self.darknet_obj_detected_cb)
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.simulation = int(rospy.get_param("/simulation"))
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
self.darknet_bboxes = []
# Simulation : If diagonal size of bounding box is more than 85px
# Simulation : Bounding Box class should be 'traffic light' with probability >= 85%
simulation_bboxes_size_threshold = 85 #px
simulation_bboxes_probability = 0.85
# Site : If diagonal size of bounding box is more than 80px
# Site : Bounding Box class should be 'traffic light' with probability >= 25%
site_bboxes_size_threshold = 40 #px
site_bboxes_probability = 0.25
if self.simulation == 1:
self.prob_thresh = simulation_bboxes_probability
self.size_thresh = simulation_bboxes_size_threshold
else:
self.prob_thresh = site_bboxes_probability
self.size_thresh = site_bboxes_size_threshold
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
# use camera cb to publish red light topic
if self.UseKnownTL == True:
self.process_bbox_and_camera_img(self.darknet_bboxes,
self.camera_image)
if self.has_new_detecting_image == True:
#self.detecting_camera_image_list_orig.append(msg)
self.detecting_camera_image = msg
self.has_new_detecting_image = False
print(" into image_cb, push a image into array")
cv_image = self.bridge.imgmsg_to_cv2(self.detecting_camera_image,
"bgr8")
image_name = "detecting_camera_image_" + str(
self.detected_img_count) + ".png"
self.detected_img_count += 1
image_path = self.images_path + image_name
cv2.imwrite(image_path, cv_image)
def darknet_detected_image_cb(self, msg):
print("into darknet_detected_image_cb")
#self.detecting_camera_image_list.append(msg)
self.darknet_detecting_camera_image = msg
self.has_new_detecting_image = True
#cv_image = self.bridge.imgmsg_to_cv2(self.darknet_detecting_camera_image, "bgr8")
#image_name = "darknet_camera_image_"+str(self.darknet_img_count)+".png"
#self.darknet_img_count += 1
#image_path = self.images_path + image_name
#cv2.imwrite(image_path, cv_image)
return
def darknet_obj_detected_cb(self, msg):
print("into darknet_obj_detected_cb")
print msg.data
return
def process_bbox_and_camera_img(self, darknet_bboxes, camera_image):
light_wp, state = self.process_traffic_lights(darknet_bboxes,
camera_image)
'''
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
return state
def darknet_bboxes_cb(self, msg):
print("into darknet_bboxes_cb")
self.darknet_bboxes = []
for bbox in msg.bounding_boxes:
if str(
bbox.Class
) == 'traffic light' and bbox.probability >= self.prob_thresh:
if math.sqrt((bbox.xmin - bbox.xmax)**2 +
(bbox.ymin - bbox.ymax)**2) >= self.size_thresh:
self.darknet_bboxes.append(bbox)
if self.simulation == 1:
#camera_image=self.detecting_camera_image_list.pop(0)
camera_image = self.darknet_detecting_camera_image
#camera_image_orig=self.detecting_camera_image_list_orig.pop(0)
camera_image_orig = self.detecting_camera_image
self.process_bbox_and_camera_img(self.darknet_bboxes,
camera_image_orig)
self.darknet_bboxes = []
else:
#camera_image=self.detecting_camera_image_list_orig.pop(0)
camera_image = self.detecting_camera_image
bbox = self.darknet_bboxes.pop(0)
cv_image = self.bridge.imgmsg_to_cv2(camera_image, "bgr8")
bb_image = cv_image[bbox.ymin:bbox.ymax, bbox.xmin:bbox.xmax]
self.light_classifier.detect_light_state(bb_image)
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
"""
#pose.position.x
#pose.position.y
#TODO implement
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
return closest_idx
def get_light_state(self, light, darknet_bboxes, camera_image):
"""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)
"""
#print "into get_light_state"
# For testing, just return the light state
if self.UseKnownTL == True:
light_state = light.state
else:
cv_image = self.bridge.imgmsg_to_cv2(camera_image, "bgr8")
light_state = self.light_classifier.get_classification(
cv_image, darknet_bboxes, self.simulation)
#print light_state
return light_state
def process_traffic_lights(self, darknet_bboxes, camera_image):
"""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)
"""
#print "into process_traffic_lights"
closest_light = None
line_wp_idx = None
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_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 waypoint index
line = stop_line_positions[i]
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
# Find closest stop line waypoint index
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, darknet_bboxes,
camera_image)
return line_wp_idx, state
#self.waypoints = None
return -1, TrafficLight.UNKNOWN