zhaoyuan/modules/lib/data_process.py

import os
import numpy as np
import pandas as pd
import yaml
from modules.lib.log_manager import LogManager


class DataPreprocessing:
    def __init__(self, data_path, config_path):
        self.data_path = data_path
        self.case_name = os.path.basename(os.path.normpath(data_path))
        self.config_path = config_path

        # 初始化配置容器
        self.vehicle_config = {}
        self.safety_config = {}
        self.comfort_config = {}
        self.efficient_config = {}
        self.function_config = {}
        self.traffic_config = {}

        # 数据容器
        self.object_df = pd.DataFrame()
        self.driver_ctrl_df = pd.DataFrame()
        self.vehicle_sys_df = pd.DataFrame()
        self.ego_data_df = pd.DataFrame()
        self.obj_data = {}
        self.ego_data = {}
        self.obj_id_list = []
        self.data_quality_level = 15

        # 加载配置并处理数据
        self._get_yaml_config()
        self._process_mode()
        self.report_info = self._get_report_info(self.ego_data)

    def _process_mode(self):
        """当前支持默认处理方式"""
        self._real_process_object_df()

    def _get_yaml_config(self):
        with open(self.config_path, 'r') as f:
            full_config = yaml.safe_load(f)

        T_threshold_config = {"T_threshold": full_config.get("T_threshold", {})}
        self.vehicle_config = full_config.get("vehicle", {})
        self.safety_config = {"safety": full_config.get("safety", {})}
        self.safety_config.update(T_threshold_config)
        self.comfort_config = {"comfort": full_config.get("comfort", {})}
        self.comfort_config.update(T_threshold_config)
        self.efficient_config = {"efficient": full_config.get("efficient", {})}
        self.efficient_config.update(T_threshold_config)
        self.function_config = {"function": full_config.get("function", {})}
        self.function_config.update(T_threshold_config)
        self.traffic_config = {"traffic": full_config.get("traffic", {})}
        self.traffic_config.update(T_threshold_config)

    def _plot_coordinate_comparison(self, ego_df):
        """绘制全局坐标系与自车坐标系的运动学量对比图"""
        try:
            import matplotlib.pyplot as plt
            from matplotlib.gridspec import GridSpec
            import os

            plt.figure(figsize=(18, 10))
            gs = GridSpec(3, 3, figure=plt.gcf())  # 改为 3 行 3 列布局

            # 1. 速度对比 - 纵向分量
            ax1 = plt.subplot(gs[0, 0])
            # ax1.plot(ego_df['simTime'], ego_df['speedX'], 'r-', label='Global Vx')
            # ax1.plot(ego_df['simTime'], ego_df['speedY'], 'g-', label='Global Vy')
            ax1.plot(ego_df['simTime'], ego_df['lon_v_vehicle'], 'b-', label='Ego Frame Lon V (2D)')
            # ax1.plot(ego_df['simTime'], ego_df['lon_v_vehicle_3D'], 'y--', label='Ego Frame Lon V (3D)')
            ax1.set_title('Longitudinal Velocity Comparison')
            ax1.set_ylabel('Velocity (m/s)')
            ax1.legend()
            ax1.grid(True)

            # 2. 速度对比 - 横向分量
            ax2 = plt.subplot(gs[0, 1])
            # ax2.plot(ego_df['simTime'], ego_df['speedX'], 'r-', label='Global Vx')
            # ax2.plot(ego_df['simTime'], ego_df['speedY'], 'g-', label='Global Vy')
            ax2.plot(ego_df['simTime'], ego_df['lat_v_vehicle'], 'b-', label='Ego Frame Lat V (2D)')
            # ax2.plot(ego_df['simTime'], ego_df['lat_v_vehicle_3D'], 'y--', label='Ego Frame Lat V (3D)')
            ax2.set_title('Lateral Velocity Comparison')
            ax2.legend()
            ax2.grid(True)

            # 3. 航向角速度（新增子图）
            ax3 = plt.subplot(gs[0, 2])
            ax3.plot(ego_df['simTime'], ego_df['speedH'], 'm-', label='Heading Rate')
            ax3.set_title('Ego Heading Rate')
            ax3.set_ylabel('Yaw Rate (deg/s)')
            ax3.grid(True)

            # 4. 加速度对比 - 纵向分量
            ax4 = plt.subplot(gs[1, 0])
            # ax4.plot(ego_df['simTime'], ego_df['accelX'], 'r-', label='Global Ax')
            # ax4.plot(ego_df['simTime'], ego_df['accelY'], 'g-', label='Global Ay')
            ax4.plot(ego_df['simTime'], ego_df['lon_acc_vehicle'], 'b-', label='Ego Frame Lon A (2D)')
            # ax4.plot(ego_df['simTime'], ego_df['lon_acc_vehicle_3D'], 'y--', label='Ego Frame Lon A (3D)')
            ax4.set_title('Longitudinal Acceleration Comparison')
            ax4.set_ylabel('Acceleration (m/s²)')
            ax4.legend()
            ax4.grid(True)

            # 5. 加速度对比 - 横向分量
            ax5 = plt.subplot(gs[1, 1])
            # ax5.plot(ego_df['simTime'], ego_df['accelX'], 'r-', label='Global Ax')
            # ax5.plot(ego_df['simTime'], ego_df['accelY'], 'g-', label='Global Ay')
            ax5.plot(ego_df['simTime'], ego_df['lon_acc_rate'], 'b-', label='Ego Frame Lat A (2D)')
            # ax5.plot(ego_df['simTime'], ego_df['lat_acc_vehicle_3D'], 'y--', label='Ego Frame Lat A (3D)')
            ax5.set_title('Lateral Acceleration Comparison')
            ax5.legend()
            ax5.grid(True)

            # 6. 航向角变化（原图，平移至右侧）
            ax6 = plt.subplot(gs[1, 2])
            ax6.plot(ego_df['simTime'], ego_df['ego_posH'], 'b-', label='Ego Heading')
            ax6.set_title('Ego Vehicle Heading')
            ax6.set_ylabel('Heading (deg)')
            ax6.grid(True)

            plt.tight_layout()

            # 保存图像到当前目录
            plot_path = os.path.join(os.getcwd(), 'coordinate_transform_comparison.png')
            plt.savefig(plot_path)
            plt.close()

            logger = LogManager().get_logger()
            logger.info(f"坐标系转换对比图已保存至: {plot_path}")

        except ImportError:
            logger = LogManager().get_logger()
            logger.warning("Matplotlib未安装，无法生成坐标系转换对比图")
        except Exception as e:
            logger = LogManager().get_logger()
            logger.error(f"生成坐标系转换对比图时出错: {e}", exc_info=True)

    def _real_process_object_df(self):
        """读取并处理对象数据（含自车）"""
        try:
            merged_csv_path = os.path.join(self.data_path, "merged_ObjState.csv")
            if not os.path.exists(merged_csv_path):
                logger = LogManager().get_logger()
                logger.error(f"文件不存在: {merged_csv_path}")
                raise FileNotFoundError(f"文件不存在: {merged_csv_path}")

            df = pd.read_csv(
                merged_csv_path,
                dtype={"simTime": float},
                on_bad_lines="skip",
                na_values=["", "NA", "null", "NaN"]
            ).drop_duplicates(subset=["simTime", "simFrame", "playerId"])

            # 列名标准化
            df.columns = [col.replace("+AF8-", "_") for col in df.columns]
            df = df.dropna(subset=["type"]).reset_index(drop=True)

            # ===== 修正航向角处理 =====
            # 确保航向角在-180到180度范围内（0度正北，90度正东）
            if 'posH' in df.columns:
                df['posH'] = df['posH'] % 360  # 规范化到0-360度
                df['posH'] = np.where(df['posH'] > 180, df['posH'] - 360, df['posH'])

            # ===== 提取自车数据 =====
            # 自车ID固定为1
            EGO_PLAYER_ID = 1
            ego_mask = df['playerId'] == EGO_PLAYER_ID

            # 确保所有车辆都有pitch和roll列
            for col in ['pitch', 'roll']:
                if col not in df.columns:
                    df[col] = 0.0

            # 提取自车姿态数据
            ego_cols = ['simTime', 'posX', 'posY', 'posH', 'pitch', 'roll']
            ego_df = df[ego_mask][ego_cols].copy()
            ego_df.columns = ['simTime'] + [f'ego_{col}' for col in ego_cols[1:]]

            # 合并自车姿态信息到主表
            df = pd.merge(df, ego_df, on='simTime', how='left')

            # ===== 计算相对位置（在自车坐标系下）=====
            # 计算全局坐标差
            dx_global = df['posX'] - df['ego_posX']
            dy_global = df['posY'] - df['ego_posY']

            # 转换为弧度并计算三角函数
            ego_heading_rad = np.deg2rad(df['ego_posH'])
            cos_h = np.cos(ego_heading_rad)
            sin_h = np.sin(ego_heading_rad)

            # 向量化计算相对坐标（在自车坐标系下）
            df['x_relative_start_dist'] = dx_global * sin_h + dy_global * cos_h
            df['y_relative_start_dist'] = -dx_global * cos_h + dy_global * sin_h

            # 自车相对位置设为0
            df.loc[ego_mask, 'x_relative_start_dist'] = 0.0
            df.loc[ego_mask, 'y_relative_start_dist'] = 0.0

            # 保存原始数据
            self.object_df = df.copy()

            # 处理运动学数据（所有车辆在自车坐标系下）
            all_processed_dfs = []
            self.obj_data = {}

            for obj_id, obj_df in df.groupby("playerId"):
                # 使用自车姿态计算运动学量
                processed_df = self._calculate_kinematics_in_ego_frame(obj_df)
                self.obj_data[obj_id] = processed_df
                all_processed_dfs.append(processed_df)

            # 合并处理后的数据
            if all_processed_dfs:
                combined_processed = pd.concat(all_processed_dfs)

                # 生成唯一键
                self.object_df['merge_key'] = self.object_df['playerId'].astype(str) + '_' + self.object_df[
                    'simTime'].astype(str)
                combined_processed['merge_key'] = combined_processed['playerId'].astype(str) + '_' + combined_processed[
                    'simTime'].astype(str)

                # 合并新列
                new_columns = [col for col in combined_processed.columns
                               if col not in self.object_df.columns or col == 'merge_key']
                self.object_df = pd.merge(
                    self.object_df,
                    combined_processed[new_columns],
                    on='merge_key',
                    how='left'
                ).drop(columns=['merge_key'])

            # 设置自车数据
            self.obj_id_list = list(self.obj_data.keys())
            self.ego_data = self.obj_data.get(EGO_PLAYER_ID, None)
            if self.ego_data is not None:
                # self._plot_coordinate_comparison(self.ego_data)
                pass
        except Exception as e:
            logger = LogManager().get_logger()
            logger.error(f"处理对象数据帧时出错: {e}", exc_info=True)
            raise

    def _calculate_kinematics_in_ego_frame(self, df):
        """计算车辆在自车坐标系下的运动学量，优化版本"""
        # 创建副本避免修改原始数据，只复制必要的列以减少内存使用
        needed_cols = ['simTime', 'ego_posH', 'ego_pitch', 'ego_roll', 'speedX', 'speedY', 'accelX', 'accelY', 'speedH']
        df = df[needed_cols + [col for col in df.columns if col not in needed_cols]].copy()

        # 确保有必要的列，使用向量化操作
        for col in ['speedZ', 'accelZ']:
            if col not in df.columns:
                df[col] = 0.0

        # 处理时间序列 - 使用向量化操作
        df.sort_values(by="simTime", inplace=True)
        df.reset_index(drop=True, inplace=True)
        df["time_diff"] = df["simTime"].diff().replace(0, np.nan).fillna(method="bfill").fillna(0.01)

        # 预计算所有三角函数值，避免重复计算
        # ===== 1. 航向角处理 =====
        ego_yaw_rad = np.deg2rad(90 - df['ego_posH'].values)  # 转换为数学坐标系
        ego_pitch_rad = np.deg2rad(df['ego_pitch'].values)
        ego_roll_rad = np.deg2rad(df['ego_roll'].values)
        yaw_rad = np.deg2rad(df['ego_posH'].values)

        # 预计算所有三角函数值
        cy_math, sy_math = np.cos(ego_yaw_rad), np.sin(ego_yaw_rad)
        cp, sp = np.cos(ego_pitch_rad), np.sin(ego_pitch_rad)
        cr, sr = np.cos(ego_roll_rad), np.sin(ego_roll_rad)
        cy, sy = np.cos(yaw_rad), np.sin(yaw_rad)

        # 获取全局速度/加速度向量
        vx = df['speedX'].values  # 东
        vy = df['speedY'].values  # 北
        vz = df['speedZ'].values  # 天
        ax = df['accelX'].values
        ay = df['accelY'].values
        az = df['accelZ'].values

        # ===== 2. 三维转换 =====
        # 构建旋转矩阵 - 使用批量计算
        r00 = cy_math * cp
        r01 = cy_math * sp * sr - sy_math * cr
        r02 = cy_math * sp * cr + sy_math * sr
        r10 = sy_math * cp
        r11 = sy_math * sp * sr + cy_math * cr
        r12 = sy_math * sp * cr - cy_math * sr
        r20 = -sp
        r21 = cp * sr
        r22 = cp * cr

        # 转换速度向量 - 使用向量化计算代替einsum
        df['lon_v_vehicle_3D'] = r00 * vx + r01 * vy + r02 * vz
        df['lat_v_vehicle_3D'] = r10 * vx + r11 * vy + r12 * vz
        df['vel_z_vehicle_3D'] = r20 * vx + r21 * vy + r22 * vz

        # 转换加速度向量 - 使用向量化计算
        df['lon_acc_vehicle_3D'] = r00 * ax + r01 * ay + r02 * az
        df['lat_acc_vehicle_3D'] = r10 * ax + r11 * ay + r12 * az
        df['acc_z_vehicle_3D'] = r20 * ax + r21 * ay + r22 * az

        # ===== 3. 二维转换（与3D水平面一致）=====
        # 使用原始航向角定义（0°=北，90°=东，顺时针为正）
        # 使用向量化计算
        df['lon_v_vehicle'] = vx * sy - vy * cy
        df['lat_v_vehicle'] = vx * cy + vy * sy
        df['lon_acc_vehicle'] = ax * sy - ay * cy
        df['lat_acc_vehicle'] = ax * cy + ay * sy

        # ===== 5. 计算其他运动学量 =====
        # 速度大小 - 使用向量化计算
        df["v"] = np.sqrt(vx ** 2 + vy ** 2)

        # 加速度变化率 - 使用向量化计算
        lat_acc_diff = np.diff(df["lat_acc_vehicle"].values, prepend=df["lat_acc_vehicle"].values[0])
        lon_acc_diff = np.diff(df["lon_acc_vehicle"].values, prepend=df["lon_acc_vehicle"].values[0])
        yawrate_diff = np.diff(df["speedH"].values, prepend=df["speedH"].values[0])
        time_diff = df["time_diff"].values

        # 使用numpy操作代替pandas操作
        df["lat_acc_rate"] = np.where(time_diff == 0, 0, lat_acc_diff / time_diff)
        df["lon_acc_rate"] = np.where(time_diff == 0, 0, lon_acc_diff / time_diff)
        df["accelH"] = np.where(time_diff == 0, 0, yawrate_diff / time_diff)

        # 替换无穷值
        df.replace([np.inf, -np.inf], np.nan, inplace=True)
        df.fillna(0, inplace=True)

        return df

    def _get_report_info(self, df):
        # 提取里程与时长
        return {
            "mileage": self._mileage_cal(df),
            "duration": self._duration_cal(df)
        }

    def _mileage_cal(self, df):
        if len(df) < 2:
            return 0.0
        if df["travelDist"].nunique() == 1:
            temp_df = df.copy()
            temp_df["time_diff"] = temp_df["simTime"].diff().fillna(0)
            temp_df["avg_speed"] = (temp_df["v"] + temp_df["v"].shift()).fillna(0) / 2
            temp_df["distance_increment"] = temp_df["avg_speed"] * temp_df["time_diff"] / 3.6
            temp_df["travelDist"] = temp_df["distance_increment"].cumsum().fillna(0)
            return round(temp_df["travelDist"].iloc[-1] - temp_df["travelDist"].iloc[0], 2)
        else:
            return round(df["travelDist"].max() - df["travelDist"].min(), 2)

    def _duration_cal(self, df):
        return df["simTime"].iloc[-1] - df["simTime"].iloc[0]