pji_desktop_script/simulation/rosbag_record.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import rosbag
import sys
import os
import signal
import subprocess
from geometry_msgs.msg import PoseWithCovarianceStamped, Twist
from move_base_msgs.msg import MoveBaseActionGoal
from datetime import datetime

class RosbagRecorder:
    def __init__(self, path):
        rospy.init_node('rosbag_recorder', anonymous=True)
        rospy.set_param('/use_sim_time', False)
        self.goal_sub = rospy.Subscriber('/move_base/goal', MoveBaseActionGoal, self.goal_callback)
        self.pose_sub = rospy.Subscriber('/amcl_pose', PoseWithCovarianceStamped, self.pose_callback)
        self.cmd_vel_sub = rospy.Subscriber('/cmd_vel', Twist, self.cmd_vel_callback)

        self.goal_position = None
        self.bag_process = None
        self.recording = False
        self.pose_position = None
        self.cmd_vel = None
        self.bag_path = path

    def goal_callback(self, msg):
        # 更新时间戳为当前时间
        msg.header.stamp = rospy.Time.now()
        self.goal_position = msg.goal.target_pose.pose.position
        rospy.loginfo("Goal received, starting rosbag recording...")
        self.start_recording()

    def pose_callback(self, msg):
        # 更新时间戳为当前时间
        msg.header.stamp = rospy.Time.now()
        self.pose_position = msg.pose.pose.position
        self.check_stop_condition()

    def cmd_vel_callback(self, msg):
        # 更新时间戳为当前时间
        # msg.header.stamp = rospy.Time.now()
        self.cmd_vel = msg
        self.check_stop_condition()

    def start_recording(self):
        if not self.recording:
            # self.bag_process = subprocess.Popen(['rosbag', 'record', '/amcl_pose', '/cmd_vel', '/some_topic1', '/some_topic2'])
            self.bag_process = subprocess.Popen(['rosbag', 'record', '-O', self.bag_path, '/imu', '/obstacle_detection', '/odom', '/sys_info'], preexec_fn=os.setsid)
            self.recording = True

    def stop_recording(self):
        if self.recording and self.bag_process:
            rospy.loginfo("Stopping rosbag recording...")
            # self.bag_process.terminate()
            # self.bag_process.wait(timeout=10)
            # os.killpg(os.getpgid(self.bag_process.pid), signal.SIGTERM)  # 首先发送SIGTERM信号
            os.killpg(os.getpgid(self.bag_process.pid), signal.SIGINT)  # 首先发送SIGINT信号
            # self.bag_process.wait(timeout=10)
            try:
                self.bag_process.wait(timeout=5)  # 等待进程终止，如果超过5秒，进行强制终止
            except subprocess.TimeoutExpired:
                rospy.logwarn("Rosbag process did not terminate, sending SIGKILL...")
                os.killpg(os.getpgid(self.bag_process.pid), signal.SIGKILL)  # 强制终止
                self.bag_process.wait()  # 再次等待确保进程已经被终止
            self.recording = False

    def check_stop_condition(self):
        if self.goal_position and self.pose_position and self.cmd_vel:
            distance = self.calculate_distance(self.pose_position, self.goal_position)
            if distance < 0.3 and self.cmd_vel.linear.x == 0 and self.cmd_vel.linear.y == 0:
                self.stop_recording()

    def calculate_distance(self, pos1, pos2):
        return ((pos1.x - pos2.x) ** 2 + (pos1.y - pos2.y) ** 2) ** 0.5

if __name__ == '__main__':
    try:
        bag_path = sys.argv[1]
        recorder = RosbagRecorder(bag_path)
        rospy.spin()
    except rospy.ROSInterruptException:
        pass