pji_desktop_script/simulation/rosbag_record.py

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

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

class RosbagRecorder:
    def __init__(self, bag_dir, count):
        # 初始化ROS节点
        rospy.init_node('rosbag_recorder', anonymous=True)
        rospy.set_param('/use_sim_time', False)        

        # 参数初始化
        self.bag_dir = bag_dir
        self.total_count = int(count)
        self.bag_count = 0
        self.recording = False
        self.goal_position = None
        self.pose_position = None
        self.cmd_vel = None
        self.bag_process = None

        # 创建锁，确保线程安全
        self.lock = threading.Lock()

        # 订阅话题
        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)

        rospy.loginfo("录制脚本初始化完成，等待终点...")

    def goal_callback(self, msg):
        with self.lock:
            if self.bag_count >= self.total_count:
                rospy.loginfo("已达到录制数量上限，退出程序。")
                self.shutdown_node()
                return

            if self.recording:
                rospy.logwarn("当前正在录制中，等待当前录制完成后再开始新的录制。")
                return

            # 记录目标位置
            self.goal_position = msg.goal.target_pose.pose.position
            rospy.loginfo(f"接收到第{self.bag_count + 1}/{self.total_count}个终点，开始录制第{self.bag_count + 1}/{self.total_count}个rosbag。")
            # 设置rosbag文件名
            msg.header.stamp = rospy.Time.now()
            bag_file_name = f"test{self.bag_count + 1:02d}.bag"
            self.bag_name = bag_file_name
            bag_file_path = os.path.join(self.bag_dir, bag_file_name)   
            self.bag_path = bag_file_path         
            # 开始录制
            self.start_recording()

    def pose_callback(self, msg):
        with self.lock:
            # msg.header.stamp = rospy.Time.now()
            self.pose_position = msg.pose.pose.position
            self.check_stop_condition()

    def cmd_vel_callback(self, msg):
        with self.lock:
            # msg.header.stamp = rospy.Time.now()
            self.cmd_vel = msg
            self.check_stop_condition()

    def check_stop_condition(self):
        if not self.recording:
            return

        if not self.goal_position or not self.pose_position or not self.cmd_vel:
            return

        # 计算距离
        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:
            rospy.loginfo(f"满足停止录制条件，停止录制第{self.bag_count + 1}个rosbag。")
            self.stop_recording()

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

    def start_recording(self):
        if self.recording:
            rospy.logwarn("已经在录制中，不启动新的录制。")
            return

        # 生成文件名，包含时间戳以避免重名
        # timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        # bag_file_name = f"test{self.bag_count + 1:02d}_{timestamp}.bag"
        # msg.header.stamp = rospy.Time.now()
        # bag_file_name = f"test{self.bag_count + 1:02d}.bag"
        # bag_file_path = os.path.join(self.bag_dir, bag_file_name)

        # # 确保目录存在
        # if not os.path.exists(self.bag_dir):
        #     try:
        #         os.makedirs(self.bag_dir)
        #         rospy.loginfo(f"创建目录：{self.bag_dir}")
        #     except Exception as e:
        #         rospy.logerr(f"无法创建目录 {self.bag_dir}：{e}")
        #         self.shutdown_node()
        #         return

        # 启动rosbag录制
        try:
            self.bag_process = subprocess.Popen(
                ['rosbag', 'record', '-o', self.bag_path, '/cmd_vel', '/move_base/goal', '/amcl_pose', '/imu', '/obstacle_detection', '/odom', '/sys_info'],
                preexec_fn=os.setsid,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            self.recording = True
            rospy.loginfo(f"开始录制rosbag：{self.bag_name}")
        except Exception as e:
            rospy.logerr(f"启动rosbag录制失败：{e}")
            self.recording = False

    def stop_recording(self):
        if not self.recording or not self.bag_process:
            rospy.logwarn("尝试停止录制，但当前没有录制进程。")
            return

        try:
            # 发送SIGTERM信号终止rosbag进程
            os.killpg(os.getpgid(self.bag_process.pid), signal.SIGTERM)
            rospy.loginfo("发送SIGTERM信号，尝试终止rosbag录制进程。")

            # 等待进程终止
            self.bag_process.wait(timeout=10)
            rospy.loginfo(f"{self.bag_name}录制进程已终止。")
        except subprocess.TimeoutExpired:
            rospy.logwarn("rosbag录制进程未能在10秒内终止，发送SIGKILL信号强制终止。")
            try:
                os.killpg(os.getpgid(self.bag_process.pid), signal.SIGKILL)
                self.bag_process.wait()
                rospy.loginfo("rosbag录制进程已被SIGKILL信号强制终止。")
            except Exception as e:
                rospy.logerr(f"强制终止rosbag录制进程失败：{e}")
        except Exception as e:
            rospy.logerr(f"停止rosbag录制进程时出错：{e}")

        # 更新状态
        self.recording = False
        self.bag_count += 1
        rospy.loginfo(f"第{self.bag_count}/{self.total_count}个rosbag录制完成。")
        # print("bag_count: ", self.bag_count)
        # print("total_count: ", self.total_count)

        # 检查是否达到录制数量
        if self.bag_count >= self.total_count:
            # rospy.loginfo("所有rosbag录制完成，退出程序。")
            self.shutdown_node()
            # rospy.loginfo("所有rosbag录制完成，退出程序。")
        else:
            rospy.loginfo("等待下一个目标点以开始新的rosbag录制。")

    def shutdown_node(self):
        rospy.loginfo("所有rosbag录制完成")
        rospy.loginfo("正在关闭ROS节点并退出程序。")
        rospy.signal_shutdown("完成所有rosbag录制。")
        rospy.sleep(1)  # 确保所有日志输出
        os._exit(0)  # 强制退出程序

if __name__ == '__main__':
    try:
        # if len(sys.argv) != 3:
        #     print("用法: rosrun your_package your_script.py <bag_directory> <count>")
        #     sys.exit(1)

        bag_folder_path = sys.argv[1]
        count = sys.argv[2]

        recorder = RosbagRecorder(bag_folder_path, count)
        rospy.spin()
    except rospy.ROSInterruptException:
        pass