planning: add traffic_light data into learning data

modules/planning/pipeline/BUILD

@@ -15,6 +15,7 @@ cc_library(
         "//modules/common/adapters:adapter_gflags",
         "//modules/common/configs:config_gflags",
         "//modules/localization/proto:localization_proto",
+        "//modules/perception/proto:perception_proto",
         "//modules/planning/common/util:util_lib",
         "//modules/planning/proto:learning_data_proto",
         "//modules/prediction/util:data_extraction",

modules/planning/pipeline/feature_generator.cc

@@ -41,6 +41,7 @@ using apollo::canbus::Chassis;
 using apollo::cyber::record::RecordMessage;
 using apollo::cyber::record::RecordReader;
 using apollo::localization::LocalizationEstimate;
+using apollo::perception::TrafficLightDetection;
 using apollo::routing::RoutingResponse;
 
 void FeatureGenerator::Init() {
@@ -76,10 +77,18 @@ void FeatureGenerator::GenerateTrajectoryLabel(
       localization_for_label.back().header().timestamp_sec());
   learning_data_frame->set_frame_num(total_learning_data_frame_num_++);
 
+  // add traffic_light
+  learning_data_frame->clear_traffic_light();
+  for (const auto& tl : traffic_lights_) {
+    auto traffic_light = learning_data_frame->add_traffic_light();
+    traffic_light->set_id(tl.first);
+    traffic_light->set_color(tl.second);
+  }
+
   // add routing
   auto features = learning_data_frame->mutable_routing_response();
   features->Clear();
-  for (const auto& lane_id : routing_lane_ids) {
+  for (const auto& lane_id : routing_lane_ids_) {
     features->add_lane_id(lane_id);
   }
 
@@ -167,22 +176,35 @@ void FeatureGenerator::OnChassis(const apollo::canbus::Chassis& chassis) {
   features->set_gear_location(chassis.gear_location());
 }
 
-void FeatureGenerator::OnRoutingResponse(
-  const apollo::routing::RoutingResponse& routing_response) {
-  if (learning_data_frame_ == nullptr) {
-    AERROR << "learning_data_frame_ pointer is nullptr";
-    return;
+void FeatureGenerator::OnTafficLightDetection(
+    const TrafficLightDetection& traffic_light_detection) {
+  // AINFO << "traffic_light_detection received at frame["
+  //      << total_learning_data_frame_num_ << "]";
+
+  traffic_lights_.clear();
+  for (int i = 0; i < traffic_light_detection.traffic_light_size(); ++i) {
+    const auto& traffic_light_id =
+        traffic_light_detection.traffic_light(i).id();
+    if (traffic_light_id.empty()) continue;
+
+    // AINFO << "  traffic_light_id[" << traffic_light_id
+    //      << "] color[" << traffic_light_detection.traffic_light(i).color() << "]";
+    traffic_lights_[traffic_light_id] =
+        traffic_light_detection.traffic_light(i).color();
   }
+}
+
 
+void FeatureGenerator::OnRoutingResponse(
+  const apollo::routing::RoutingResponse& routing_response) {
   AINFO << "routing_response received at frame["
         << total_learning_data_frame_num_ << "]";
-
-  routing_lane_ids.clear();
+  routing_lane_ids_.clear();
   for (int i = 0; i < routing_response.road_size(); ++i) {
     for (int j = 0; j < routing_response.road(i).passage_size(); ++j) {
       for (int k = 0; k < routing_response.road(i).passage(j).segment_size();
           ++k) {
-        routing_lane_ids.push_back(
+        routing_lane_ids_.push_back(
             routing_response.road(i).passage(j).segment(k).id());
       }
     }
@@ -213,6 +235,11 @@ void FeatureGenerator::ProcessOfflineData(const std::string& record_filename) {
       if (routing_response.ParseFromString(message.content)) {
         OnRoutingResponse(routing_response);
       }
+    } else if (message.channel_name == FLAGS_traffic_light_detection_topic) {
+      TrafficLightDetection traffic_light_detection;
+      if (traffic_light_detection.ParseFromString(message.content)) {
+        OnTafficLightDetection(traffic_light_detection);
+      }
     }
   }
 }

modules/planning/pipeline/feature_generator.h

@@ -18,12 +18,14 @@
 #include <list>
 #include <string>
 #include <vector>
+#include <unordered_map>
 
 #include "cyber/common/file.h"
 #include "modules/canbus/proto/chassis.pb.h"
 #include "modules/localization/proto/localization.pb.h"
-#include "modules/routing/proto/routing.pb.h"
+#include "modules/perception/proto/traffic_light_detection.pb.h"
 #include "modules/planning/proto/learning_data.pb.h"
+#include "modules/routing/proto/routing.pb.h"
 
 namespace apollo {
 namespace planning {
@@ -40,6 +42,8 @@ class FeatureGenerator {
   void OnLocalization(const apollo::localization::LocalizationEstimate& le);
   void OnRoutingResponse(
       const apollo::routing::RoutingResponse& routing_response);
+  void OnTafficLightDetection(
+      const apollo::perception::TrafficLightDetection& traffic_light_detection);
 
   void WriteOutLearningData(const LearningData& learning_data,
                             const std::string& file_name);
@@ -55,7 +59,9 @@ class FeatureGenerator {
   std::list<apollo::localization::LocalizationEstimate>
       localization_for_label_;
   int total_learning_data_frame_num_ = 0;
-  std::vector<std::string> routing_lane_ids;
+  std::vector<std::string> routing_lane_ids_;
+  std::unordered_map<std::string, apollo::perception::TrafficLight::Color>
+        traffic_lights_;
 };
 
 }  // namespace planning

modules/planning/proto/learning_data.proto

@@ -93,7 +93,7 @@ message RoutingResponseFeature {
 }
 
 message TrafficLightFeature {
-    optional string light_id = 1;
+    optional string id = 1;
     optional apollo.perception.TrafficLight.Color color = 2;
 }