Browse Source

补充新增指标需要的lane和road的coords信息

XGJ_zhaoyuan 1 tháng trước cách đây
mục cha
commit
6e0de80f24
1 tập tin đã thay đổi với 54 bổ sung20 xóa
  1. 54 20
      core/processors/built_in/lst.py

+ 54 - 20
core/processors/built_in/lst.py

@@ -32,6 +32,8 @@ OUTPUT_CSV_LANEMAP = "LaneMap.csv"
 OUTPUT_CSV_EGOMAP = "EgoMap.csv"
 OUTPUT_CSV_FUNCTION = "Function.csv"
 ROADMARK_CSV = "RoadMark.csv"
+HD_LANE_CSV = "hd_lane.csv"
+HD_LINK_CSV = "hd_link.csv"
 
 
 # --- Configuration Class ---
@@ -365,8 +367,7 @@ class ZipCSVProcessor:
                 processed_data.append(record)
 
             if processed_data:
-                df_final = pd.DataFrame(processed_data)[output_columns].iloc[::4].reset_index(
-                    drop=True)  # Ensure column order
+                df_final = pd.DataFrame(processed_data)[output_columns].iloc[::4].reset_index(drop=True)  # Ensure column order
                 df_final['simFrame'] = np.arange(1, len(df_final) + 1)
                 df_final.to_csv(output_path, index=False, encoding="utf-8")
                 print(f"Successfully wrote GNSS data to {output_path}")
@@ -427,8 +428,8 @@ class ZipCSVProcessor:
             df_obj = df_obj.reindex(columns=final_columns).iloc[::4]
 
             # Reindex simFrame of ego and obj
-            df_ego['simFrame'] = np.arange(1, len(df_ego) + 1)
-            df_obj['simFrame'] = np.arange(1, len(df_obj) + 1)
+            df_ego['simFrame'] = np.arange(1, len(df_ego)+1)
+            df_obj['simFrame'] = np.arange(1, len(df_obj)+1)
 
             # Merge EGO and OBJ dataframes
             df_merged = pd.concat([df_ego, df_obj], ignore_index=True)
@@ -581,9 +582,7 @@ class ZipCSVProcessor:
         # df_vehicle.dropna(subset=[col for col in required_pos if col in df_vehicle.columns], inplace=True)
 
         try:
-            df_vehicle["simTime"] = np.round(np.linspace(df_raw["simTime"].tolist()[0] + 28800,
-                                                         df_raw["simTime"].tolist()[0] + 28800 + 0.01 * (
-                                                             len(df_vehicle)), len(df_vehicle)), 2)
+            df_vehicle["simTime"] = np.round(np.linspace(df_raw["simTime"].tolist()[0]+28800, df_raw["simTime"].tolist()[0]+28800 + 0.01*(len(df_vehicle)), len(df_vehicle)), 2)
             df_vehicle["simFrame"] = np.arange(1, len(df_vehicle) + 1)
             df_vehicle["playerId"] = int(player_id)
             df_vehicle['playerId'] = pd.to_numeric(df_vehicle['playerId']).astype(int)
@@ -784,6 +783,7 @@ class PolynomialCurvatureFitting:
             return [
                 {
                     "projection": (np.nan, np.nan),
+                    "lightMask": 0,
                     "curvHor": np.nan,
                     "curvHorDot": np.nan,
                     "laneOffset": np.nan,
@@ -797,6 +797,7 @@ class PolynomialCurvatureFitting:
         for x_target, y_target in points:
             result = {  # Default result
                 "projection": (np.nan, np.nan),
+                "lightMask": 0,
                 "curvHor": np.nan,
                 "curvHorDot": np.nan,
                 "laneOffset": np.nan,
@@ -831,6 +832,7 @@ class PolynomialCurvatureFitting:
 
             result = {
                 "projection": (round(proj_x, 6), round(proj_y, 6)),
+                "lightMask": 0,
                 "curvHor": round(curv_hor, 6),
                 "curvHorDot": round(curv_hor_dot, 6),
                 "laneOffset": round(min_distance, 6),
@@ -1067,6 +1069,7 @@ class FinalDataProcessor:
 
             indices = df_object[df_object["playerId"] == 1].index
             if len(indices) == len(curvHor_values):
+                df_object.loc[indices, "lightMask"] = 0
                 df_object.loc[indices, "curvHor"] = curvHor_values
                 df_object.loc[indices, "curvHorDot"] = curvature_change_value
                 df_object.loc[indices, "laneOffset"] = min_distance
@@ -1169,8 +1172,8 @@ class FinalDataProcessor:
                 df_ego = pd.read_csv(egomap_path, dtype={"simTime": float})
                 ego_column = ['posX', 'posY', 'posH']
                 ego_new_column = ['posX_map', 'posY_map', 'posH_map']
-                df_ego = df_ego.rename(columns=dict(zip(ego_column, ego_new_column)))
-                # 删除 simFrame 列,因为使用主数据的 simFrame
+                df_ego = df_ego.rename(columns = dict(zip(ego_column, ego_new_column)))
+            # 删除 simFrame 列,因为使用主数据的 simFrame
                 if 'simFrame' in df_ego.columns:
                     df_ego = df_ego.drop(columns=['simFrame'])
 
@@ -1203,7 +1206,7 @@ class FinalDataProcessor:
                     else:
                         print("警告: 合并 EgoMap 后找不到 simTime 列!")
 
-                df_merged = df_merged.drop(columns=['posX_map', 'posY_map', 'posH_map'])
+                df_merged = df_merged.drop(columns = ['posX_map', 'posY_map', 'posH_map'])
 
                 print("EgoMap data merged.")
             except Exception as e:
@@ -1214,6 +1217,37 @@ class FinalDataProcessor:
         # 先处理可能的列名重复问题
         df_merged = clean_duplicate_columns(df_merged)
 
+        # --- 合并hd_lane.csv,hd_road.csv ---
+        current_file_path = os.path.abspath(__file__)
+        root_lane_csv_path1 = os.path.dirname(current_file_path)
+        root_lane_csv_path2 = os.path.dirname(root_lane_csv_path1)
+        root_lane_csv_path3 = os.path.dirname(root_lane_csv_path2)
+        root_lane_csv_path4 = os.path.dirname(root_lane_csv_path3)
+        lane_path = os.path.join(root_lane_csv_path4, "_internal")
+        data_path = os.path.join(lane_path, "data_map")
+        lane_csv_path = os.path.join(data_path, "hd_lane.csv")
+        road_csv_path = os.path.join(data_path, "hd_link.csv")
+        df_lane = pd.read_csv(lane_csv_path)
+        column_to_read = ['link_type', 'link_coords']
+
+        df_road = pd.read_csv(road_csv_path, usecols = column_to_read)
+        df_road["simFrame"] = np.arange(1, len(df_road) + 1, 1)
+        # df_road = df_road.rename(columns={'link_type': 'road_type'})
+
+
+        df_merged = pd.merge(
+            df_merged,
+            df_lane,
+            on='lane_id',
+            how = 'left'
+        )
+        df_merged = pd.merge(
+            df_merged,
+            df_road,
+            on='simFrame',
+            how = 'left'
+        )
+
         # --- 合并 Traffic ---
         traffic_path = self.output_dir / "Traffic.csv"
         if traffic_path.exists() and traffic_path.stat().st_size > 0:
@@ -1264,12 +1298,12 @@ class FinalDataProcessor:
                     7: 'S',  # 南左转
                     8: 'W',  # 西左转
                     9: 'N',  # 北左转
-                    10: 'E',  # 东左转
-                    11: 'N',  # 北行人
-                    12: 'E',  # 东行人
-                    13: 'S',  # 南右转
-                    14: 'W',  # 西右转
-                    15: 'N',  # 北右转
+                    10: 'E', # 东左转
+                    11: 'N', # 北行人
+                    12: 'E', # 东行人
+                    13: 'S', # 南右转
+                    14: 'W', # 西右转
+                    15: 'N', # 北右转
                     16: 'E'  # 东右转
                 }
 
@@ -1339,11 +1373,10 @@ class FinalDataProcessor:
                     if phase_direction == row['vehicle_direction']:
                         # 查找该方向的所有红绿灯 ID
                         relevant_ids = [tid for tid, direction in trafficlight_to_direction.items()
-                                        if direction == phase_direction]
+                                       if direction == phase_direction]
 
                         # 如果 trafficlight_id 在 EgoMap 中且方向匹配
-                        if trafficlight_col in row and not pd.isna(row[trafficlight_col]) and row[
-                            trafficlight_col] in relevant_ids:
+                        if trafficlight_col in row and not pd.isna(row[trafficlight_col]) and row[trafficlight_col] in relevant_ids:
                             return row[trafficlight_col]
 
                     return np.nan
@@ -1499,6 +1532,7 @@ class FinalDataProcessor:
 
         return df_merged
 
+
     def _process_trafficlight_data(self) -> pd.DataFrame:
         """Processes traffic light JSON data if available."""
         # Check if json_path is provided and exists
@@ -1593,7 +1627,6 @@ class FinalDataProcessor:
             print(f"Unexpected error processing traffic light data: {e}")
             return pd.DataFrame()
 
-
 # --- Rosbag Processing ---
 class RosbagProcessor:
     """Extracts data from Rosbag files within a ZIP archive."""
@@ -1656,6 +1689,7 @@ class RosbagProcessor:
         if not target_rostopic:
             print("Skipping Rosbag processing as no target topic was identified.")
 
+
         with tempfile.TemporaryDirectory() as tmp_dir_str:
             tmp_dir = Path(tmp_dir_str)
             bag_files_extracted = []