【python 自定义保存ros下bag包的某话题图像jpg】

功能

列出当前文件夹内的所有.bag文件 bag包的数据类型sensor_msgs/CompressedImage
让用户选择一个bag进行进一步操作
列出所选.bag文件中的所有话题，让用户选择一个话题
创建以话题名命名的文件夹的
解析该话题下的图像数据，将图像以时间戳命名保存到文件夹内

# -*- coding: utf-8 -*-
import os
import rosbag
from cv_bridge import CvBridge
import cv2
from sensor_msgs.msg import CompressedImage

def list_bag_files():
    """列出当前目录下的所有.bag文件"""
    return [file for file in os.listdir('.') if file.endswith('.bag')]

def select_bag_file(bag_files):
    """让用户从列表中选择一个.bag文件"""
    for i, file in enumerate(bag_files):
        print("{}{}. {}".format(i+1, ".", file))
    selection = int(input("请选择要解析的bag包的编号: ")) - 1
    return bag_files[selection]

def list_topics_in_bag(bag_file):
    """列出指定bag文件中的所有话题"""
    bag = rosbag.Bag(bag_file, 'r')
    topics = bag.get_type_and_topic_info()[1].keys()
    bag.close()
    return topics

def select_topic(topics):
    """让用户从列表中选择一个话题"""
    for i, topic in enumerate(topics):
        print("{}{}. {}".format(i+1, ".", topic))
    selection = int(input("请选择要解析的话题的编号: ")) - 1
    return list(topics)[selection]

def create_topic_folder(topic_name):
    """创建以话题名称为名的文件夹，如果文件夹不存在"""
    folder_name = topic_name.replace('/', '_')[1:]  # 移除话题名称前的斜杠并替换剩余的斜杠为下划线
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)
    return folder_name

def extract_images_from_bag(bag_file, topic):
    """从选定的bag包和话题中提取图像数据"""
    bag = rosbag.Bag(bag_file, 'r')
    bridge = CvBridge()
    topic_folder = create_topic_folder(topic)  # 创建存储图像的文件夹

    for _, msg, t in bag.read_messages(topics=[topic]):
        if hasattr(msg, 'header'):
            timestamp_secs = msg.header.stamp.secs
            timestamp_nsecs = msg.header.stamp.nsecs
            timestamp_msecs = int(timestamp_nsecs / 1000000)
            datetime_format = "{}.{:03}".format(timestamp_secs, timestamp_msecs)
            image_name = 'img_{}.jpg'.format(datetime_format)
            image_path = os.path.join(topic_folder, image_name)  # 构建图像文件的完整路径

            cv_image = bridge.compressed_imgmsg_to_cv2(msg, desired_encoding="passthrough")
            cv2.imwrite(image_path, cv_image)
            print("图像已保存至：{}".format(image_path))

    bag.close()
    print("完成。图像已提取至'{}'文件夹。".format(topic_folder))

if __name__ == '__main__':
    bag_files = list_bag_files()
    if not bag_files:
        print("当前目录下没有找到.bag文件。")
    else:
        selected_bag_file = select_bag_file(bag_files)
        topics = list_topics_in_bag(selected_bag_file)
        if not topics:
            print("选定的bag包中没有找到话题。")
        else:
            selected_topic = select_topic(topics)
            extract_images_from_bag(selected_bag_file, selected_topic)