Skip to content

Instantly share code, notes, and snippets.

Created February 10, 2022 03:54
Show Gist options
  • Save rancheng/11b13ec6c2182781af7fdc11ed7c3c18 to your computer and use it in GitHub Desktop.
Save rancheng/11b13ec6c2182781af7fdc11ed7c3c18 to your computer and use it in GitHub Desktop.
bag file to image files in tum dataset and realsense dataset
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2021
"""Extract images from a rosbag and synchronize the sequences
Please run this script with /usr/bin/python2.7 instead of python3
and don't forget to source the ros environment: source /opt/ros/melodic/setup.bash
import os
import argparse
import numpy as np
import cv2
import copy
import rosbag
import shutil
from tqdm import tqdm
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
def matching_time_indices(stamps_1, stamps_2,
max_diff = 0.1,
offset_2 = 0.0):
Searches for the best matching timestamps of two lists of timestamps
and returns the list indices of the best matches.
:param stamps_1: first vector of timestamps (numpy array)
:param stamps_2: second vector of timestamps (numpy array)
:param max_diff: max. allowed absolute time difference
:param offset_2: optional time offset to be applied to stamps_2
:return: 2 lists of the matching timestamp indices (stamps_1, stamps_2)
matching_indices_1 = []
matching_indices_2 = []
stamps_2 = copy.deepcopy(stamps_2)
stamps_2 += offset_2
for index_1, stamp_1 in enumerate(stamps_1):
diffs = np.abs(stamps_2 - stamp_1)
index_2 = int(np.argmin(diffs))
if diffs[index_2] <= max_diff:
return matching_indices_1, matching_indices_2
def main():
"""Extract a folder of images from a rosbag.
global seq_base_folder_sync
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("--tum", help="convert to tum mode.", dest="tum", type=bool, default=False)
parser.add_argument("--image_topic", nargs='+', help="usage: --image_topic /image/data /depth/data")
args = parser.parse_args()
print "Extract images from %s on topic %s into %s" % (args.bag_file,
args.image_topic, args.output_dir)
bag = rosbag.Bag(args.bag_file, "r")
bridge = CvBridge()
color_count = 0
depth_count = 0
color_stamps = {}
depth_stamps = {} # use depth stamps as master and color as aligner for time stamp sync
if not os.path.exists(args.output_dir):
# you can use topics= [args.image_topic] to specify your topic name /device_0/sensor_0/Depth_0/image/data /device_0/sensor_1/Color_0/image/data
# for topic, msg, t in bag.read_messages(topics=["/camera/color/image_raw", "/camera/aligned_depth_to_color/image_raw"]):
if args.image_topic:
img_topics = args.image_topic
img_topics = ["/camera/color/image_raw", "/camera/aligned_depth_to_color/image_raw"]
for topic, msg, t in bag.read_messages(
cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
output_fname = ""
if not 'depth' in topic.lower() and 'color' in topic.lower():
cv_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
color_directory_path = os.path.join(args.output_dir, "color")
if not os.path.exists(color_directory_path):
output_fname = os.path.join(color_directory_path, "%06i.png" % color_count)
print "Wrote color image %i" % color_count
color_stamps[msg.header.stamp.to_sec()] = output_fname
color_count += 1
if 'depth' in topic.lower():
depth_directory_path = os.path.join(args.output_dir, "depth")
if not os.path.exists(depth_directory_path):
output_fname = os.path.join(depth_directory_path, "%06i.png" % depth_count)
print "Wrote depth image %i" % depth_count
depth_stamps[msg.header.stamp.to_sec()] = output_fname
depth_count += 1
cv2.imwrite(output_fname, cv_img)
print("starting time sync...")
# start the time sync:
depth_stamps_t = np.fromiter(depth_stamps.iterkeys(), dtype=float)
color_stamps_t = np.fromiter(color_stamps.iterkeys(), dtype=float)
# find the matching indices between depth and color time stamps
matching_indices_1, matching_indices_2 = matching_time_indices(depth_stamps_t, color_stamps_t)
# len(matching_indices_1) == len(matching_indices_2)
matched_depth_stamps_t = []
seq_base_folder_sync = args.output_dir + "_sync"
if args.tum:
seq_base_folder_sync += "_tum"
seq_depth_folder_sync = os.path.join(seq_base_folder_sync, "depth")
if args.tum:
seq_color_folder_sync = os.path.join(seq_base_folder_sync, "rgb")
seq_color_folder_sync = os.path.join(seq_base_folder_sync, "color")
if not os.path.exists(seq_base_folder_sync):
for depth_iter_idx, dpt_index in tqdm(enumerate(matching_indices_1)):
# obtain the file path of the matched indices
depth_file_path = depth_stamps[depth_stamps_t[dpt_index]]
color_file_path = color_stamps[color_stamps_t[matching_indices_2[depth_iter_idx]]]
# save the depth and color sync files to the new folder with new indices
if args.tum:
# print "Converting to TUM dataset"
shutil.copy(depth_file_path, os.path.join(seq_depth_folder_sync, "%1.6f.png" % depth_stamps_t[dpt_index]))
shutil.copy(color_file_path, os.path.join(seq_color_folder_sync, "%1.6f.png" % depth_stamps_t[dpt_index]))
# print "Converting to RealSense Dataset"
shutil.copy(depth_file_path, os.path.join(seq_depth_folder_sync, "%06i.png" % depth_iter_idx))
shutil.copy(color_file_path, os.path.join(seq_color_folder_sync, "%06i.png" % depth_iter_idx))
# record the time stamp
matched_depth_stamps_t.append([depth_iter_idx, depth_stamps_t[dpt_index]])
matched_depth_stamps_t = np.array(matched_depth_stamps_t)
print("end time sync.")
fmt = '%06d', '%1.6f'
np.savetxt(os.path.join(seq_base_folder_sync, "timestamps.txt"), matched_depth_stamps_t, fmt)
print("timestamps.txt saved to %s" % os.path.join(seq_base_folder_sync, "timestamps.txt"))
if __name__ == '__main__':
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment