使用 Python 从 ROS Bag 中提取图像：详解与实现

在机器人应用中，ROS (Robot Operating System) 是一个常见的框架。ROS Bag（rosbag）是 ROS 中用于记录和回放数据流（例如传感器数据、话题消息等）的一种强大工具。有时，我们需要将存储在 rosbag 文件中的图像数据提取并保存为图像文件以便进一步分析或处理。本文将介绍如何编写一个 Python 脚本，从 rosbag 文件中提取图像并保存为 PNG 文件。

功能概述

该脚本的主要功能包括：

1. 读取指定的 ROS Bag 文件。
2. 从指定的话题（Topic）中提取图像数据。
3. 使用 OpenCV 将图像保存为 PNG 格式文件。
4. 提供灵活的命令行参数，支持不同的输入文件、输出目录和话题。

注意：

如果ROS Bag中的图像数据为sensor_msgs/CompressedImage通过以下方式先转换为sensor_msgs/Image，重新录制一个Bag

rosrun image_transport republish compressed in:=/camera/color/image_raw raw out:=/camera/color/image_raw

脚本实现

下面是完整的 Python 脚本代码：

1. 普通RGB（sensor_msgs/Image）图像

#!/usr/bin/env python3
import argparse
import cv2
import os
import rosbag
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

def extract_images_from_bag(bag_file, output_dir, image_topic):
    # 打开rosbag文件
    bag = rosbag.Bag(bag_file, 'r')
    bridge = CvBridge()
    count = 0

    # 读取指定话题的消息
    for topic, msg, t in bag.read_messages(topics=[image_topic]):
        try:
            # 将ROS消息转换为OpenCV图像
            cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            print(f"Error converting image: {e}")
            continue

        # 保存为图像文件
        image_filename = os.path.join(output_dir, f"frame_{count:06d}.png")
        cv2.imwrite(image_filename, cv_image)
        count += 1
        print(f"Image {count} saved to {image_filename}")

    # 关闭rosbag文件
    bag.close()
    print(f"Processed {count} images.")

def main():
    # 使用argparse处理命令行参数
    parser = argparse.ArgumentParser(description="Extract images from a rosbag and save them as files.")
    parser.add_argument("bag_file", help="The rosbag file to extract images from")
    parser.add_argument("output_dir", help="Directory to save the extracted images")
    parser.add_argument("image_topic", help="Image topic to subscribe to")

    args = parser.parse_args()

    # 确保输出目录存在
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    # 从rosbag中提取图像
    extract_images_from_bag(args.bag_file, args.output_dir, args.image_topic)

if __name__ == '__main__':
    main()

2. 8UC3红外图像

#!/usr/bin/env python3
import argparse
import cv2
import os
import rosbag
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

def extract_images_from_bag(bag_file, output_dir, image_topic):
    # 打开rosbag文件
    bag = rosbag.Bag(bag_file, 'r')
    bridge = CvBridge()
    count = 0

    # 读取指定话题的消息
    for topic, msg, t in bag.read_messages(topics=[image_topic]):
        try:
            # 检查图像的编码格式
            if msg.encoding == '8UC3':
                # 直接转换为OpenCV图像
                cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
            else:
                # 转换为指定的颜色编码（例如 'bgr8'）
                cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            print(f"Error converting image: {e}")
            continue

        # 保存为图像文件
        image_filename = os.path.join(output_dir, f"frame_{count:06d}.png")
        cv2.imwrite(image_filename, cv_image)
        count += 1
        print(f"Image {count} saved to {image_filename}")

    # 关闭rosbag文件
    bag.close()
    print(f"Processed {count} images.")

def main():
    # 使用argparse处理命令行参数
    parser = argparse.ArgumentParser(description="Extract images from a rosbag and save them as files.")
    parser.add_argument("bag_file", help="The rosbag file to extract images from")
    parser.add_argument("output_dir", help="Directory to save the extracted images")
    parser.add_argument("image_topic", help="Image topic to subscribe to")

    args = parser.parse_args()

    # 确保输出目录存在
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    # 从rosbag中提取图像
    extract_images_from_bag(args.bag_file, args.output_dir, args.image_topic)

if __name__ == '__main__':
    main()

脚本讲解

1. 依赖库

该脚本依赖以下库：

安装依赖库

在运行脚本前，需要确保已安装这些依赖项。以下是安装命令：

pip install opencv-python
sudo apt install python3-rosbag python3-cv-bridge

脚本功能详解

2. 主要功能模块

2.1 从 ROS Bag 中提取图像

extract_images_from_bag 函数是脚本的核心部分，主要完成以下任务：

1. 读取 Bag 文件
   使用 rosbag.Bag 打开指定的 ROS Bag 文件以便提取数据。

2. 遍历消息
   使用 bag.read_messages 遍历指定话题中的所有消息。

3. 转换图像
   借助 cv_bridge 将 ROS 格式的图像消息（sensor_msgs/Image）转换为 OpenCV 格式的图像数据。

4. 保存图像
   使用 cv2.imwrite 将提取的图像保存为 PNG 文件，文件名格式为 frame_000001.png、frame_000002.png 等。

2.2 命令行参数解析

该脚本使用 argparse 支持灵活的命令行参数配置，支持以下参数：

通过这些参数，用户可以灵活配置脚本，处理不同的输入文件、输出路径和图像来源话题。

2.3 确保目录存在

为了确保图像可以正确保存，脚本在保存图像之前会检查目标输出目录是否存在：

3. 运行脚本

使用以下命令运行脚本：

python3 extract_images.py <bag_file> <output_dir> <image_topic>

示例

假设

运行脚本的命令

在终端中运行以下命令：

python3 extract_images.py data.bag output /camera/image_raw

输出结果

脚本运行后将执行以下操作：

1. 从指定的话题中提取图像数据：

脚本会读取 ROS Bag 文件中的图像数据，并从指定的话题（例如 /camera/image_raw）中提取图像消息。

2. 保存图像到指定的输出目录：

提取的图像会以 PNG 格式保存在 output 目录中，用户可以通过该目录查看保存的图像文件。

3. 文件命名格式：

图像文件将按照顺序命名为 frame_000001.png、frame_000002.png 等。例如，如果提取了 100 张图像，则会生成文件 frame_000001.png 到 frame_000100.png。

4. 终端输出进度：

每提取一张图像，脚本会在终端输出其保存路径。完成后，还会显示总共提取并保存了多少张图像。

示例输出

Image 1 saved to output/frame_000001.png
Image 2 saved to output/frame_000002.png
Image 3 saved to output/frame_000003.png
...
Processed 100 images.

脚本运行完成后，用户可以在 output 目录中找到所有提取的图像文件。

作者：点云兔子