zhaoyuan/modules/lib/data_process.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
##################################################################
#
# Copyright (c) 2024 CICV, Inc. All Rights Reserved
#
##################################################################
"""
@Authors:           zhanghaiwen(zhanghaiwen@china-icv.cn)
@Data:              2024/10/17
@Last Modified:     2024/10/17
@Summary:           Evaluation functions
"""

import os

import numpy as np
import pandas as pd

import yaml

from modules.lib.log_manager import LogManager


# from lib import log  # 确保这个路径是正确的，或者调整它
# logger = None  # 初始化为 None


class DataPreprocessing:
    def __init__(self, data_path, config_path):
        # Initialize paths and data containers
        # self.logger = log.get_logger()

        self.data_path = data_path
        self.case_name = os.path.basename(os.path.normpath(data_path))

        self.config_path = config_path

        # Initialize DataFrames
        self.object_df = pd.DataFrame()
        self.driver_ctrl_df = pd.DataFrame()
        self.vehicle_sys_df = pd.DataFrame()
        self.ego_data_df = pd.DataFrame()

        # Environment data
        self.lane_info_df = pd.DataFrame()
        self.road_mark_df = pd.DataFrame()
        self.road_pos_df = pd.DataFrame()
        self.traffic_light_df = pd.DataFrame()
        self.traffic_signal_df = pd.DataFrame()

        self.vehicle_config = {}
        self.safety_config = {}
        self.comfort_config = {}
        self.efficient_config = {}
        self.function_config = {}
        self.traffic_config = {}

        # Initialize data for later processing
        self.obj_data = {}
        self.ego_data = {}
        self.obj_id_list = []

        # Data quality level
        self.data_quality_level = 15

        # Process mode and prepare report information
        self._process_mode()
        self._get_yaml_config()
        self.report_info = self._get_report_info(self.obj_data.get(1, pd.DataFrame()))

    def _process_mode(self):
        """Handle different processing modes."""
        self._real_process_object_df()

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

        modules = ["vehicle", "T_threshold", "safety", "comfort", "efficient", "function", "traffic"]

        # 1. 初始化 vehicle_config（不涉及 T_threshold 合并）
        self.vehicle_config = full_config[modules[0]]

        # 2. 定义 T_threshold_config（封装为字典）
        T_threshold_config = {"T_threshold": full_config[modules[1]]}

        # 3. 统一处理需要合并 T_threshold 的模块
        # 3.1 safety_config
        self.safety_config = {"safety": full_config[modules[2]]}
        self.safety_config.update(T_threshold_config)

        # 3.2 comfort_config
        self.comfort_config = {"comfort": full_config[modules[3]]}
        self.comfort_config.update(T_threshold_config)

        # 3.3 efficient_config
        self.efficient_config = {"efficient": full_config[modules[4]]}
        self.efficient_config.update(T_threshold_config)

        # 3.4 function_config
        self.function_config = {"function": full_config[modules[5]]}
        self.function_config.update(T_threshold_config)

        # 3.5 traffic_config
        self.traffic_config = {"traffic": full_config[modules[6]]}
        self.traffic_config.update(T_threshold_config)

    @staticmethod
    def cal_velocity(lat_v, lon_v):
        """Calculate resultant velocity from lateral and longitudinal components."""
        return np.sqrt(lat_v ** 2 + lon_v ** 2)

    def _real_process_object_df(self):
        """Process the object DataFrame."""
        try:
            # 读取 CSV 文件
            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}")

            self.object_df = pd.read_csv(
                merged_csv_path,
                dtype={"simTime": float},
                engine="python",
                on_bad_lines="skip",  # 自动跳过异常行
                na_values=["", "NA", "null", "NaN"]  # 明确处理缺失值
            ).drop_duplicates(subset=["simTime", "simFrame", "playerId"])
            self.object_df.columns = [col.replace("+AF8-", "_") for col in self.object_df.columns]

            data = self.object_df.copy()

            # 使用向量化操作计算速度和加速度，提高性能
            data["lat_v"] = data["speedY"] * 1
            data["lon_v"] = data["speedX"] * 1
            # 使用向量化操作代替 apply
            data["v"] = np.sqrt(data["lat_v"] ** 2 + data["lon_v"] ** 2)

            # 计算加速度分量
            data["lat_acc"] = data["accelY"] * 1
            data["lon_acc"] = data["accelX"] * 1
            # 使用向量化操作代替 apply
            data["accel"] = np.sqrt(data["lat_acc"] ** 2 + data["lon_acc"] ** 2)

            # Drop rows with missing 'type' and reset index
            data = data.dropna(subset=["type"])
            data.reset_index(drop=True, inplace=True)
            self.object_df = data.copy()

            # Calculate respective parameters for each object
            for obj_id, obj_data in data.groupby("playerId"):
                self.obj_data[obj_id] = self._calculate_object_parameters(obj_data)

            # Get object id list
            EGO_PLAYER_ID = 1
            self.obj_id_list = list(self.obj_data.keys())
            self.ego_data = self.obj_data[EGO_PLAYER_ID]

            # 添加这一行：处理自车数据，进行坐标系转换
            self.ego_data = self.process_ego_data(self.ego_data)

        except Exception as e:
            logger = LogManager().get_logger()
            logger.error(f"处理对象数据帧时出错: {e}", exc_info=True)
            raise

    def _calculate_object_parameters(self, obj_data):
        """Calculate additional parameters for a single object."""
        obj_data = obj_data.copy()

        # 确保数据按时间排序
        obj_data = obj_data.sort_values(by="simTime").reset_index(drop=True)

        obj_data["time_diff"] = obj_data["simTime"].diff()

        # 处理可能的零时间差
        zero_time_diff = obj_data["time_diff"] == 0
        if zero_time_diff.any():
            logger = LogManager().get_logger()
            logger.warning(f"检测到零时间差: {sum(zero_time_diff)} 行")
            # 将零时间差替换为最小非零时间差或一个小的默认值
            min_non_zero = obj_data.loc[~zero_time_diff, "time_diff"].min() if (~zero_time_diff).any() else 0.01
            obj_data.loc[zero_time_diff, "time_diff"] = min_non_zero

        obj_data["lat_acc_diff"] = obj_data["lat_acc"].diff()
        obj_data["lon_acc_diff"] = obj_data["lon_acc"].diff()
        obj_data["yawrate_diff"] = obj_data["speedH"].diff()

        obj_data["lat_acc_roc"] = (
                obj_data["lat_acc_diff"] / obj_data["time_diff"]
        ).replace([np.inf, -np.inf], [9999, -9999])
        obj_data["lon_acc_roc"] = (
                obj_data["lon_acc_diff"] / obj_data["time_diff"]
        ).replace([np.inf, -np.inf], [9999, -9999])
        obj_data["accelH"] = (
                obj_data["yawrate_diff"] / obj_data["time_diff"]
        ).replace([np.inf, -np.inf], [9999, -9999])

        return obj_data

    def _get_driver_ctrl_data(self, df):
        """
        Process and get driver control information.

        Args:
            df: A DataFrame containing driver control data.

        Returns:
            A dictionary of driver control info.
        """
        driver_ctrl_data = {
            "time_list": df["simTime"].round(2).tolist(),
            "frame_list": df["simFrame"].tolist(),
            "brakePedal_list": (
                (df["brakePedal"] * 100).tolist()
                if df["brakePedal"].max() < 1
                else df["brakePedal"].tolist()
            ),
            "throttlePedal_list": (
                (df["throttlePedal"] * 100).tolist()
                if df["throttlePedal"].max() < 1
                else df["throttlePedal"].tolist()
            ),
            "steeringWheel_list": df["steeringWheel"].tolist(),
        }
        return driver_ctrl_data

    def _get_report_info(self, df):
        """Extract report information from the DataFrame."""
        mileage = self._mileage_cal(df)
        duration = self._duration_cal(df)
        return {"mileage": mileage, "duration": duration}

    def _mileage_cal(self, df):
        """Calculate mileage based on the driving data."""
        if len(df) < 2:
            return 0.0  # 数据不足，无法计算里程

        if df["travelDist"].nunique() == 1:
            # 创建临时DataFrame进行计算，避免修改原始数据
            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)

            mileage = round(temp_df["travelDist"].iloc[-1] - temp_df["travelDist"].iloc[0], 2)
            return mileage
        else:
            # 如果travelDist已经有多个值，直接计算最大值和最小值的差
            return round(df["travelDist"].max() - df["travelDist"].min(), 2)
        return 0.0  # Return 0 if travelDist is not valid

    def _duration_cal(self, df):
        """Calculate duration of the driving data."""
        return df["simTime"].iloc[-1] - df["simTime"].iloc[0]

    def process_ego_data(self, ego_data):
        """处理自车数据：将东北天（ENU）坐标系下的速度/加速度转换为车辆坐标系（考虑yaw, pitch, roll）"""
        '''
        原字段	新字段名	描述
        a_x_body	lon_acc_vehicle	车辆坐标系下的纵向加速度
        a_y_body	lat_acc_vehicle	车辆坐标系下的横向加速度
        a_z_body	acc_z_vehicle	车辆坐标系下的垂向加速度
        v_x_body	lon_v_vehicle	车辆坐标系下的纵向速度
        v_y_body	lat_v_vehicle	车辆坐标系下的横向速度
        v_z_body	vel_z_vehicle	车辆坐标系下的垂向速度
        posH	heading_rad	航向角（弧度）
        pitch_rad	pitch_rad	俯仰角（弧度）
        roll_rad	roll_rad	横滚角（弧度）
        '''
        logger = LogManager().get_logger()

        if ego_data is None or len(ego_data) == 0:
            logger.warning("自车数据为空，无法进行坐标系转换")
            return ego_data

        ego_data = ego_data.copy()
        for col in ['speedZ', 'accelZ']:
            if col not in ego_data.columns:
                ego_data[col] = 0.0
                logger.warning(f"自车数据中缺少列 '{col}'，已将其填充为 0.0")

        # 角度转为弧度（修正 posH 表示正北为 0° => 车辆朝正东为 0°）
        ego_data['yaw_rad'] = np.deg2rad(90 - ego_data['posH'])
        ego_data['pitch_rad'] = np.deg2rad(ego_data.get('pitch', 0))
        ego_data['roll_rad'] = np.deg2rad(ego_data.get('roll', 0))

        # 预计算三角函数（向量化）
        cy = np.cos(ego_data['yaw_rad'])
        sy = np.sin(ego_data['yaw_rad'])
        cp = np.cos(ego_data['pitch_rad'])
        sp = np.sin(ego_data['pitch_rad'])
        cr = np.cos(ego_data['roll_rad'])
        sr = np.sin(ego_data['roll_rad'])

        # === 加速度（ENU → 车辆坐标系） ===
        ego_data['lon_acc_vehicle'] = (ego_data['accelX'] * (cy * cp) +
                                       ego_data['accelY'] * (cy * sp * sr - sy * cr) +
                                       ego_data['accelZ'] * (cy * sp * cr + sy * sr))

        ego_data['lat_acc_vehicle'] = (ego_data['accelX'] * (sy * cp) +
                                       ego_data['accelY'] * (sy * sp * sr + cy * cr) +
                                       ego_data['accelZ'] * (sy * sp * cr - cy * sr))

        ego_data['acc_z_vehicle'] = (ego_data['accelX'] * (-sp) +
                                     ego_data['accelY'] * (cp * sr) +
                                     ego_data['accelZ'] * (cp * cr))

        # === 速度（ENU → 车辆坐标系） ===
        ego_data['lon_v_vehicle'] = (ego_data['speedX'] * (cy * cp) +
                                     ego_data['speedY'] * (cy * sp * sr - sy * cr) +
                                     ego_data['speedZ'] * (cy * sp * cr + sy * sr))

        ego_data['lat_v_vehicle'] = (ego_data['speedX'] * (sy * cp) +
                                     ego_data['speedY'] * (sy * sp * sr + cy * cr) +
                                     ego_data['speedZ'] * (sy * sp * cr - cy * sr))

        ego_data['vel_z_vehicle'] = (ego_data['speedX'] * (-sp) +
                                     ego_data['speedY'] * (cp * sr) +
                                     ego_data['speedZ'] * (cp * cr))

        logger.info("完成车辆坐标系转换（考虑yaw/pitch/roll）")
        return ego_data