A simple script to extract synced topics at a fixed rate from a rosbag

process_rosbag_sync.py

"""
File: process_rosbag_sync.py
Author: Dhruv Shah (shah@cs.berkeley.edu)
Python Version: 3.8.13

A simple script to extract rostopics at a fixed (reduced) rate from a rosbag.
"""

import pickle as pkl
import pandas
import numpy as np
import io
import os
import matplotlib.pyplot as plt
import rosbag
from PIL import Image

## change this as needed
robotname = 'spot'
idx = 0 
datadir = "/media/tempSSD/datadir"
csvfile = 'data_{}.csv'.format(robotname)
imtopic = '/image_raw/compressed'
odomtopic = '/localization'
rate = 4.

df = pandas.read_csv(csvfile, header=None)
filenames = df.values
filenames = [name[0] for name in filenames]
bagname = filenames[idx]

# open bag
bag = rosbag.Bag(os.path.join(datadir, bagname))

# check if bag has both topics
if not bag.get_message_count(imtopic) or not bag.get_message_count(odomtopic):
    print("bag does not have both topics")
    exit()
# get image and odom data from bags
synced_imdata = []
synced_odomdata = []

# get start time of bag in seconds
currtime = bag.get_start_time()
starttime = currtime

curr_imdata = None
curr_odomdata = None
times = []

for topic, msg, t in bag.read_messages(topics=[imtopic, odomtopic]):
    if topic == imtopic:
        curr_imdata = msg.data
    elif topic == odomtopic:
        curr_odomdata = msg.pose
    if (t.to_sec() - currtime) >= 1. / rate:
        if curr_imdata is not None and curr_odomdata is not None:
            synced_imdata.append(curr_imdata)
            synced_odomdata.append(curr_odomdata)
        else: 
            print("missing data")
        currtime = t.to_sec()
        times.append(currtime - starttime)

# decompress images
synced_images = [Image.open(io.BytesIO(imdata)) for imdata in synced_imdata]

# resize image to X by 120, preserving aspect ratio
synced_images_resized = [im.resize((int(120 * im.size[0] / im.size[1]), 120)) for im in synced_images]

# process odom data
trajdata = {}
trajdata['position'] = np.array([[odom.x, odom.y] for odom in synced_odomdata])
trajdata['yaw'] = np.array([odom.theta for odom in synced_odomdata])
savedir = '/nfs/kun2/users/dshah/multirobot/scand/auto'
bagdir = robotname + '_' + bagname.split('.')[0]
if os.path.exists(os.path.join(savedir, bagdir)):
    print("bag already processed")
    exit()
else:
    os.mkdir(os.path.join(savedir, bagdir))

    # save images as jpegs
    for i, im in enumerate(synced_images):
        im.save(os.path.join(savedir, bagdir, str(i) + '.jpg'))
    
    # save trajdata as a pickle file
    with open(os.path.join(savedir, bagdir, 'traj_data.pkl'), 'wb') as f:
        pkl.dump(trajdata, f)