Rosbag Extractor in Python

import os
from argparse import ArgumentParser
import rosbag
import numpy as np
import cv2

parser = ArgumentParser()
parser.add_argument("-d", dest = "out_dir", help='output directory')
parser.add_argument("bag_name", help='the bag name')

args = parser.parse_args()

outdir = args.out_dir
bagname = args.bag_name

if outdir is None:
    outdir = bagname.split('.')[0]

if not os.path.exists(outdir):
    os.makedirs(outdir)

target_cameras = [ ('/gmsl_camera/port_a/cam_0/image_raw/compressed', 'cam60_left'), 
                   ('/gmsl_camera/port_a/cam_1/image_raw/compressed', 'cam60_middle'),
                   ('/gmsl_camera/port_a/cam_2/image_raw/compressed', 'cam60_right'),
                   ('/gmsl_camera/port_b/cam_1/image_raw/compressed', 'cam30_middle') ]

bag = rosbag.Bag(bagname)

def extractJPG(t_cam):
    count = 0 
    cam_path = os.path.join(outdir, t_cam[1])
    print t_cam[0], t_cam[1]

    if not os.path.exists(cam_path):
        os.makedirs(cam_path)
    else:
        print "Already generated!"
        return

    for topic, msg, t in bag.read_messages(topics=[t_cam[0]]):
        byte_array = np.fromstring(msg.data, np.uint8)
        img = cv2.imdecode(byte_array, cv2.IMREAD_COLOR)
        name = os.path.join(cam_path, "img" + "{:0>4d}".format(count) + ".jpg")
        cv2.imwrite(name, img)
        count += 1    


for target_camera in target_cameras:
    extractJPG(target_camera)

print outdir, " Done!"

Last updated

Was this helpful?