diff --git a/modules/prediction/container/adc_trajectory/adc_trajectory_container.cc b/modules/prediction/container/adc_trajectory/adc_trajectory_container.cc
index 0031513e76a52b9a389f4cb6a3673f6727d0a85b..1ee1727a8f594550f810461302a214c203735b0a 100644
--- a/modules/prediction/container/adc_trajectory/adc_trajectory_container.cc
+++ b/modules/prediction/container/adc_trajectory/adc_trajectory_container.cc
@@ -35,6 +35,7 @@ using apollo::planning::ADCTrajectory;
 
 void ADCTrajectoryContainer::Insert(
     const ::google::protobuf::Message& message) {
+  reference_line_lane_ids_.clear();
   adc_trajectory_.CopyFrom(dynamic_cast<const ADCTrajectory&>(message));
   if (!IsProtected()) {
     junction_polygon_ = Polygon2d{};
@@ -55,8 +56,8 @@ bool ADCTrajectoryContainer::IsPointInJunction(const PathPoint& point) const {
     on_virtual_lane = map->IsVirtualLane(point.lane_id());
   }
   if (!on_virtual_lane) {
-    on_virtual_lane = map->OnVirtualLane({point.x(), point.y()},
-                                         FLAGS_virtual_lane_radius);
+    on_virtual_lane =
+        map->OnVirtualLane({point.x(), point.y()}, FLAGS_virtual_lane_radius);
   }
   return in_polygon && on_virtual_lane;
 }
@@ -67,44 +68,46 @@ bool ADCTrajectoryContainer::IsProtected() const {
 }
 
 Polygon2d ADCTrajectoryContainer::GetJunctionPolygon() {
-  int num_point = adc_trajectory_.trajectory_point_size();
-  if (num_point == 0) {
-    return Polygon2d{};
-  }
   std::shared_ptr<const JunctionInfo> junction_info(nullptr);
-  double prev_s = 0.0;
-  for (int i = 0; i < num_point; ++i) {
+  for (int i = 0; i < adc_trajectory_.trajectory_point_size(); ++i) {
     double s = adc_trajectory_.trajectory_point(i).path_point().s();
-    if (i > 0 && std::abs(s - prev_s) < FLAGS_junction_search_radius) {
-      continue;
-    }
-    if (s > FLAGS_adc_trajectory_search_length) {
-      break;
+    std::string lane_id =
+        adc_trajectory_.trajectory_point(i).path_point().lane_id();
+
+    // Find junctions
+    if (junction_info == nullptr && s < FLAGS_adc_trajectory_search_length) {
+      double x = adc_trajectory_.trajectory_point(i).path_point().x();
+      double y = adc_trajectory_.trajectory_point(i).path_point().y();
+      std::vector<std::shared_ptr<const JunctionInfo>> junctions =
+          PredictionMap::instance()->GetJunctions({x, y},
+                                                  FLAGS_junction_search_radius);
+      if (!junctions.empty() && junctions.front() != nullptr) {
+        junction_info = junctions.front();
+      }
     }
-    prev_s = s;
-    double x = adc_trajectory_.trajectory_point(i).path_point().x();
-    double y = adc_trajectory_.trajectory_point(i).path_point().y();
-    std::vector<std::shared_ptr<const JunctionInfo>> junctions =
-        PredictionMap::instance()->GetJunctions({x, y},
-                                                FLAGS_junction_search_radius);
-    if (junctions.empty()) {
-      continue;
-    } else {
-      junction_info = junctions.front();
-      break;
+
+    // Insert reference lane ids
+    if (reference_line_lane_ids_.empty() ||
+        lane_id != reference_line_lane_ids_.back()) {
+      reference_line_lane_ids_.emplace_back(lane_id);
     }
   }
-  if (junction_info == nullptr) {
-    return Polygon2d{};
-  }
 
-  const apollo::hdmap::Polygon& junction_polygon =
-      junction_info->junction().polygon();
-  std::vector<Vec2d> vertices;
-  for (const auto& point : junction_polygon.point()) {
-    vertices.emplace_back(point.x(), point.y());
+  if (junction_info != nullptr) {
+    std::vector<Vec2d> vertices;
+    for (const auto& point : junction_info->junction().polygon().point()) {
+      vertices.emplace_back(point.x(), point.y());
+    }
+    if (vertices.size() >= 3) {
+      return Polygon2d{vertices};
+    }
   }
-  return Polygon2d{vertices};
+  return Polygon2d{};
+}
+
+const std::vector<std::string>&
+ADCTrajectoryContainer::get_reference_line_lane_ids() {
+  return reference_line_lane_ids_;
 }
 
 }  // namespace prediction
diff --git a/modules/prediction/container/adc_trajectory/adc_trajectory_container.h b/modules/prediction/container/adc_trajectory/adc_trajectory_container.h
index 7d38982328b03b4d08266a5170fcdfaef2a2ffd4..f293557ac925f5551a71ef68a82fe3a18d0ca86c 100644
--- a/modules/prediction/container/adc_trajectory/adc_trajectory_container.h
+++ b/modules/prediction/container/adc_trajectory/adc_trajectory_container.h
@@ -24,7 +24,6 @@
 
 #include <memory>
 #include <string>
-#include <unordered_set>
 #include <vector>
 
 #include "modules/perception/proto/perception_obstacle.pb.h"
@@ -68,13 +67,19 @@ class ADCTrajectoryContainer : public Container {
    */
   bool IsPointInJunction(const apollo::common::PathPoint& point) const;
 
+  /**
+   * @brief Get reference line lane ids
+   * @return A vector of lane ids
+   */
+  const std::vector<std::string>& get_reference_line_lane_ids();
+
  private:
   apollo::common::math::Polygon2d GetJunctionPolygon();
 
  private:
   apollo::planning::ADCTrajectory adc_trajectory_;
   apollo::common::math::Polygon2d junction_polygon_;
-  std::unordered_set<std::string> reference_line_lane_ids_;
+  std::vector<std::string> reference_line_lane_ids_;
 };
 
 }  // namespace prediction