ROS笔记之提取源Bag文件中指定的Topic到新的Bag文件中(python版)

extractSelectedTopic.py

# coding=utf-8
import rosbag
import subprocess, yaml
import sys


def getSummaryInfo(bag_path):
    info_strs = []
    info_dict = yaml.safe_load(
        subprocess.Popen(['rosbag', 'info', '--yaml', bag_path],
                         stdout=subprocess.PIPE).communicate()[0])
    end_timestamp = float(info_dict['end'])
    duration = float(info_dict['duration'])
    start_timestamp = end_timestamp - duration

    start_time_str = "Start timestamp:" + str(start_timestamp) + " s"
    end_time_str = "End timestamp:" + str(end_timestamp) + " s"
    duration_str = "Duration:" + str(duration) + " s"

    print("-" * 100)
    print("Summary Info:")
    print(start_time_str)
    print(end_time_str)
    print(duration_str)
    info_strs.append("-" * 100 + "\n")
    info_strs.append("Summary Info:\n")
    info_strs.append(start_time_str + "\n")
    info_strs.append(end_time_str + "\n")
    info_strs.append(duration_str + "\n")

    print("Topic list:")
    info_strs.append("Topic list:\n")

    topic_list = info_dict['topics']
    img_topic_list = []
    for i in range(len(topic_list)):
        if topic_list[i]['type'].__contains__("Image"):
            img_topic_list.append(topic_list[i])
        topic_info_item = "No." + str(i + 1) + ": Name: " + topic_list[i][
            'topic'] + ", \t\tType: " + topic_list[i][
                'type'] + ", \t\tMessages:" + str(
                    topic_list[i]['messages']) + ", \t\tFramerate:" + str(
                        int(topic_list[i]['messages'] / duration))
        print(topic_info_item)
        info_strs.append(topic_info_item + "\n")
    print("-" * 100)
    info_strs.append("-" * 100 + "\n")
    return info_dict, topic_list, img_topic_list, start_timestamp, end_timestamp, duration, info_strs


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("You should input one bag filepath at least.")
        exit()
    bag_path = sys.argv[1]
    out_path = bag_path[:bag_path.rfind(".")] + "_extracted.bag"

    info_dict, topic_list, img_topic_list, start_timestamp, end_timestamp, duration, info_strs = getSummaryInfo(
        bag_path)

    print('topic所在index,用', '隔开,如:9,5')
    input_indices = input(
        "Input selected topic indices to output(separated by ',' ):\n")

    processed_indices = input_indices.split(',')
    for i in range(len(processed_indices)):
        processed_indices[i] = int(processed_indices[i])
        if processed_indices[i] < 1:
            processed_indices[i] = 0
        elif processed_indices[i] > len(topic_list):
            processed_indices[i] = len(topic_list) - 1
        else:
            processed_indices[i] = processed_indices[i] - 1

    processed_indices = list(set(processed_indices))
    processed_indices.sort()

    print("Selected topics:")
    for i in range(len(processed_indices)):
        print("\t" + topic_list[processed_indices[i]]['topic'])

    bag_out = rosbag.Bag(out_path, 'w')

    with rosbag.Bag(bag_path, 'r') as bag:
        for topic, msg, t in bag.read_messages():
            for i in range(len(processed_indices)):
                if topic == topic_list[processed_indices[i]]['topic']:
                    bag_out.write(topic, msg, t)
                    # print(str(t.to_sec()) + "/" + str(end_timestamp))

    bag_out.close()
    print("Processed the whole bag, saved at: " + out_path)