Python解析rosbag数据

前言

使用Python的rosbags库从rosbag1文件中提取点云和图片数据。

导入及函数准备

难点主要是解析msg的data数据，注意以下几点：

1. 所有的点是一个铺平的ndarray，每个元素代表一个字节的数据
2. 每msg.point_step个字节合起来代表一个点
3. 在一个点/周期内，存放着该点的所有msg.fields的数据
4. 每个field占几个字节，都在msg.fields里有说明

import struct
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image
from rosbags.highlevel import AnyReader


def generate_pcd_header(
    fields,
    sizes,
    types,
    n_points
):
    headers = (
        "# .PCD v0.7 - Point Cloud Data file format\n"
        "VERSION 0.7\n"
        f"FIELDS {' '.join(fields)}\n"
        f"SIZE {' '.join(map(str, sizes))}\n"
        f"TYPE {' '.join(types)}\n"
        f"COUNT {' '.join(['1'] * len(fields))}\n"
        f"WIDTH {n_points}\n"
        "HEIGHT 1\n"
        "VIEWPOINT 0 0 0 1 0 0 0\n"
        f"POINTS {n_points}\n"
        "DATA binary\n"
    )

    return headers

def parse_msg_pcd(
    msg,
    needed_fields = ('x', 'y', 'z', 'intensity')
):
    datatype_byte_sizes = {
        1: 1,  # INT8
        2: 1,  # UINT8
        3: 2,  # INT16
        4: 2,  # UINT16
        5: 4,  # INT32
        6: 4,  # UINT32
        7: 4,  # FLOAT32
        8: 8   # FLOAT64
    }
    
    datatype_unpack_formats = {
        1: 'b',  # INT8
        2: 'B',  # UINT8
        3: 'h',  # INT16
        4: 'H',  # UINT16
        5: 'i',  # INT32
        6: 'I',  # UINT32
        7: 'f',  # FLOAT32
        8: 'd'   # FLOAT64
    }

    datatype_names = {
        1: 'INT8',  # INT8
        2: 'UINT8',  # UINT8
        3: 'INT16',  # INT16
        4: 'UINT16',  # UINT16
        5: 'INT32',  # INT32
        6: 'UINT32',  # UINT32
        7: 'FLOAT32',  # FLOAT32
        8: 'FLOAT64'   # FLOAT64
    }
    
    # 提取需要的字段的信息
    n_points = msg.width
    fields = [
        x for x in msg.fields
        if x.name in needed_fields
    ]

    # 用于生成文件头
    sizes = []
    types = []
    for x in fields:
        sizes.append(datatype_byte_sizes[x.datatype])
        types.append(datatype_names[x.datatype])
    
    header = generate_pcd_header(
        needed_fields,
        sizes,
        [x[0] for x in types],
        n_points
    )
    
    # 用于从原始数据中取点
    fields_map = {
        x.name: (
            x.offset, 
            datatype_byte_sizes[x.datatype], 
            datatype_unpack_formats[x.datatype]
        ) 
        for x in fields
    }
    
    # 提取每个点的信息
    points = []
    point_step = msg.point_step
    data = msg.data
    for i in range(n_points):
        # 计算当前点的数据在 msg.data 数组中的起始索引
        index = i * point_step
        
        # 解析每个点
        single_point = []
        for k, (offset, size, format) in fields_map.items():
            bytes_data = data[index + offset : index + offset + size]
            single_point.append(
                struct.unpack(format, bytes_data)[0]
            )
        
        points.append(single_point)
        
    pc_df = pd.DataFrame(points, columns=needed_fields)

    # 转换类型
    for i, x in enumerate(needed_fields):
        pc_df[x] = pc_df[x].astype(types[i].lower())

    return header, pc_df

解析PCD

筛出所有代表点云的message，然后用上面的函数解析即可。

origin_path = Path(r'xxx.bag')
target_path = Path(r"xxx")

reader = AnyReader([origin_path])
reader.open()

# pcd
pcds = [x for x in reader.connections if x.topic == "/lidar_0/vls128/pcl2"]

# 记录pcd名字，用于和图片名对齐，可以无视
names = []
for lid_connection, lid_timestamp, lid_rawdata in reader.messages(connections=pcds):
    lid_msg = reader.deserialize(lid_rawdata, lid_connection.msgtype)
    print(lid_msg.frame_id)
    header, pc_df = parse_msg_pcd(
        lid_msg,
        needed_fields=['x', 'y', 'z', "intensity"]
    )
    pcd_data = header.encode() + pc_df.to_records(index=False).tobytes()

    stamp = lid_msg.header.stamp
    pcd_output = target_path / "lidar_point_cloud_0" / f"{stamp.sec + stamp.nanosec / 1e9}.pcd"
    pcd_output.parent.mkdir(parents=True, exist_ok=True)
    pcd_output.write_bytes(pcd_data)

    names.append(stamp.sec + stamp.nanosec / 1e9)

解析图片

图片，我这里遇到的非常简单，可以直接一个语句搞定Image.open(io.BytesIO(img_msg.data))

...

imgs = [x for x in reader.connections if x.topic == "/pylon_camera_node_0/image_color/compressed"]
for img_connection, img_timestamp, img_rawdata in reader.messages(connections=imgs):
    if not names:
        break

    cur_name = names[0] # 用于对齐，可以无视
    
    img_msg = reader.deserialize(img_rawdata, img_connection.msgtype)
    image_pil = Image.open(io.BytesIO(img_msg.data))
    stamp = img_msg.header.stamp
    if stamp.sec + stamp.nanosec / 1e9 >= float(cur_name): # if语句用于对齐，可以无视
        img_output = target_path / "camera_image_0" / f"{cur_name}.jpg"
        img_output.parent.mkdir(parents=True, exist_ok=True)
        image_pil.save(img_output)
        names.pop(0)
        
reader.close()