rosbag_to_images.py

"""
Created on Mon Dec 11 14:02:33 2017
@author: merwan
"""

import os
import argparse
import cv2
import rosbag
import sys

from sensor_msgs.msg import Image
from cv_bridge import CvBridge

class Tee(object):
    def __init__(self, *files):
        self.files = files
    def write(self, obj):
        for f in self.files:
            f.write(obj)
            f.flush() # If you want the output to be visible immediately
    def flush(self) :
        for f in self.files:
            f.flush()
            
original = sys.stdout


def extract_info_bag():
    parser = argparse.ArgumentParser(description="Extract images from a ROS bag.")
    parser.add_argument("bag_file", help="Input ROS bag.")
    parser.add_argument("output_dir", help="Output directory.")
    parser.add_argument("image_topic", help="Image topic.")

    args = parser.parse_args()

    print("Extract images from %s on topic %s into %s" % (args.bag_file,
                                                          args.image_topic, args.output_dir))
	
    # Where to write the results
    file2wtime = open(args.output_dir+'/data_timestamp.txt','w')

    bag = rosbag.Bag(args.bag_file, "r")
    bridge = CvBridge()
    count = 0
    for topic, msg, t in bag.read_messages(topics=[args.image_topic]):
        cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
        hros = msg.header
        timeh = hros.stamp.secs*1e9 + hros.stamp.nsecs

        # write in a file                    
        sys.stdout = Tee(sys.stdout, file2wtime)
        print("%06i" % count, 
        print("%0.0f" %  timeh, 
        print(9+"pic_%06i.png" % count)

        cv2.imwrite(os.path.join(args.output_dir, "pic_%06i.png" % count), cv_img)
        sys.stdout = original  
        print("Wrote image %i" % count)

        count += 1

    bag.close()

    return

if __name__ == '__main__':
    extract_info_bag()