PJI_outdoor_robot_evaluate_real/pjirebot_evaluate_csv.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Dec 25 17:35:56 2023

@author: dell
"""

# coding: utf-8
# !/usr/bin/env python2
import os
import rosbag
import csv
import math

import sys
import time
from datetime import datetime
import rospy


def parsebag(input_dir, output_dir):
    dic_EgoState = ['Time', 'simTime', 'simFrame', 'posX', 'posY', 'posZ', 'speedX', 'speedY', 'speedZ', 'accelX',
                    'accelY', 'accelZ', 'dimX', 'dimY', 'dimZ', 'obstacle', 'traveledDist']
    # dic_DriverCtrl=['Time','simTime','simFrame','tarspeedX','tarspeedY','tarspeedZ','tardimX','tardimY','tardimZ']

    EgoState_file = open(output_dir + "/" + "EgoState_pji.csv", 'w')
    # DriverCtrl_file = open(output_dir + "/"+"DriverCtrl_pji.csv", 'w')
    writer_EgoState = csv.writer(EgoState_file)
    writer_EgoState.writerow(dic_EgoState)
    # writer_DriverCtrl = csv.writer(DriverCtrl_file)
    # writer_DriverCtrl.writerow(dic_DriverCtrl)

    frame_max = sys.maxsize
    count = 1
    with rosbag.Bag(input_dir, 'r') as bag:
        odom_flag = False
        first_message_time = None
        Frame_imu = 1
        Frame_cmd_vel = 1
        obstacle_state = 0
        cur_mileage = ''

        for topic, msg, t in bag.read_messages():  # t代表时间

            if first_message_time is None:
                first_message_time = t
                first_message_time = rospy.Time.to_sec(first_message_time)
                first_message_time = datetime.fromtimestamp(first_message_time)

            if topic == "/obstacle_detection":
                obstacle_state = msg.data
                # print(msg.data)

            if topic == "/sys_info":
                cur_mileage = msg.cur_mileage

            if topic == "/odom":
                odom_flag = True
                posX = msg.pose.pose.position.x
                posY = msg.pose.pose.position.y
                posZ = msg.pose.pose.position.z
                speedX = msg.twist.twist.linear.x * 3.6
                speedY = msg.twist.twist.linear.y * 3.6
                speedZ = msg.twist.twist.linear.z * 3.6
                dimX = msg.twist.twist.angular.x
                dimY = msg.twist.twist.angular.y
                dimZ = msg.twist.twist.angular.z

            if topic == "/imu":

                if odom_flag:
                    accelX = msg.linear_acceleration.x
                    accelY = msg.linear_acceleration.y
                    accelZ = msg.linear_acceleration.z
                    timestamp = rospy.Time.to_sec(t)
                    date_time_imu = datetime.fromtimestamp(timestamp)
                    simtime_imu = (date_time_imu - first_message_time).total_seconds()

                    message_EgoState = [date_time_imu, simtime_imu, Frame_imu, posX, posY, posZ, speedX, speedY, speedZ,
                                        accelX, accelY, accelZ, dimX, dimY, dimZ, obstacle_state, cur_mileage]

                    writer_EgoState.writerow(message_EgoState)
                    Frame_imu += 1
                else:
                    print('6666')

            '''        
            if topic =='/cmd_vel':
                timestamp = rospy.Time.to_sec(t)
                date_time_cmd_vel = datetime.fromtimestamp(timestamp)
                simtime_cmd_vel=(date_time_cmd_vel-first_message_time).total_seconds()
                message_DriverCtrl =[date_time_cmd_vel,simtime_cmd_vel,Frame_cmd_vel,msg.linear.x*3.6,msg.linear.y*3.6,msg.linear.z*3.6,
                                     msg.angular.x,msg.angular.y,msg.angular.z]
                
                                     
                writer_DriverCtrl.writerow(message_DriverCtrl)
                Frame_cmd_vel+=1
            '''

        EgoState_file.close()
        # DriverCtrl_file.close()


if __name__ == "__main__":
    #    print("Input the path of file...."+'\n')

    # input_dir = sys.argv[1]
    input_dir = './bag/2024-04-22-08-48-30_obstacledetection_30.bag'
    bagname = input_dir.split('/')[-1].split('.')[0]
    # output_dir = sys.argv[2]
    output_dir = './'
    output_dir = os.path.join(output_dir, bagname)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    parsebag(input_dir, output_dir)

    '''
    try:
        parsebag(input_dir, output_dir)
        print('successfully analysis '+input_dir)
    except Exception as e:
        print(e)
    '''