def _read_info(self):
"""Read some metadata from inside this rosbag."""
from rosbag import Bag, ROSBagUnindexedException, ROSBagException
try:
b = Bag(self.path)
except ROSBagUnindexedException:
b = Bag(self.path, allow_unindexed=True)
print('Reindexing', self.filename)
b.reindex()
try:
duration = b.get_end_time() - b.get_start_time()
except ROSBagException:
duration = 0
return {
'n_messages': b.get_message_count(),
'duration': duration,
'topic_list': b.get_type_and_topic_info()[1].keys()
}
def main():
parser = argparse.ArgumentParser(
description=("Extracts grayscale and event images from a ROS bag and "
"saves them as TFRecords for training in TensorFlow."))
parser.add_argument("--bag", dest="bag",
help="Path to ROS bag.",
required=True)
parser.add_argument("--prefix", dest="prefix",
help="Output file prefix.",
required=True)
parser.add_argument("--output_folder", dest="output_folder",
help="Output folder.",
required=True)
parser.add_argument("--max_aug", dest="max_aug",
help="Maximum number of images to combine for augmentation.",
type=int,
default=6)
parser.add_argument("--n_skip", dest="n_skip",
help="Maximum number of images to combine for augmentation.",
type=int,
default=1)
parser.add_argument("--start_time", dest="start_time",
help="Time to start in the bag.",
type=float,
default=0.0)
parser.add_argument("--end_time", dest="end_time",
help="Time to end in the bag.",
type=float,
default=-1.0)
parser.add_argument("--save_rgb_images", default=True,
const=False, nargs="?")
parser.add_argument("--debug", default=False,
const=True, nargs="?")
parser.add_argument("--whitelist_imageids_txt",
type=str,
default=None)
args = parser.parse_args()
bridge = CvBridge()
n_msgs = 0
left_start_event_offset = 0
right_start_event_offset = 0
left_event_image_iter = 0
right_event_image_iter = 0
left_image_iter = 0
right_image_iter = 0
first_left_image_time = -1
first_right_image_time = -1
left_events = []
right_events = []
left_images = []
right_images = []
left_image_times = []
right_image_times = []
left_event_count_images = []
left_event_time_images = []
left_event_image_times = []
right_event_count_images = []
right_event_time_images = []
right_event_image_times = []
left_image_event_start_idx = []
left_image_event_end_idx = []
right_image_event_start_idx = []
right_image_event_end_idx = []
whitelist_imageids = None
if args.whitelist_imageids_txt is not None:
with open(args.whitelist_imageids_txt, 'r') as fp:
whitelist_imageids = fp.read().splitlines()
whitelist_imageids = [int(l) for l in whitelist_imageids]
cols = 346
rows = 260
print("Processing bag")
bag = Bag(args.bag)
h5_left, h5_right = None, None
if args.debug:
import h5py
h5_file = h5py.File(args.bag[:-len("bag")]+"hdf5")
h5_left = h5_file['davis']['left']['events']
h5_right = h5_file['davis']['right']['events']
options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.GZIP)
left_tf_writer = tf.python_io.TFRecordWriter(
os.path.join(args.output_folder, args.prefix, "left_event_images.tfrecord"),
options=options)
right_tf_writer = tf.python_io.TFRecordWriter(
os.path.join(args.output_folder, args.prefix, "right_event_images.tfrecord"),
options=options)
# Get actual time for the start of the bag.
t_start = bag.get_start_time()
t_start_ros = rospy.Time(t_start)
# Set the time at which the bag reading should end.
if args.end_time == -1.0:
t_end = bag.get_end_time()
else:
t_end = t_start + args.end_time
for topic, msg, t in bag.read_messages(
topics=['/davis/left/image_raw',
'/davis/right/image_raw',
'/davis/left/events',
'/davis/right/events'],
# **** MOD **** #
# NOTE: we always start reading from the start in order
# to count the number of events that have to be
# discarded in the HDF5 file
# start_time=rospy.Time(args.start_time + t_start),
end_time=rospy.Time(t_end)):
# Check to make sure we're working with stereo messages.
if not ('left' in topic or 'right' in topic):
print('ERROR: topic {} does not contain left or right, is this stereo?'
'If not, you will need to modify the topic names in the code.'.
format(topic))
return
n_msgs += 1
if n_msgs % 500 == 0:
print("Processed {} msgs, {} images, time is {}.".format(n_msgs,
left_event_image_iter,
t.to_sec() - t_start))
isLeft = 'left' in topic
# **** MOD **** # /*start
# If we are still not reading the part
# we are interested in, we just count
# the number of events
if t.to_sec() < args.start_time + t_start:
if 'events' in topic and msg.events:
if isLeft:
left_start_event_offset += len(msg.events)
else:
right_start_event_offset += len(msg.events)
continue # read the next msg
# **** MOD **** # end*/
if 'image' in topic:
width = msg.width
height = msg.height
if width != cols or height != rows:
print("Image dimensions are not what we expected: set: ({} {}) vs got:({} {})"
.format(cols, rows, width, height))
return
time = msg.header.stamp
image = np.asarray(bridge.imgmsg_to_cv2(msg, msg.encoding))
image = np.reshape(image, (height, width))
if isLeft:
if whitelist_imageids is None or left_image_iter in whitelist_imageids:
if args.save_rgb_images:
cv2.imwrite(os.path.join(args.output_folder,
args.prefix,
"left_image{:05d}.png".format(left_image_iter)),
image)
if left_image_iter > 0:
left_image_times.append(time)
else:
first_left_image_time = time
left_event_image_times.append(time.to_sec())
left_image_iter += 1
else:
if whitelist_imageids is None or right_image_iter in whitelist_imageids:
if args.save_rgb_images:
cv2.imwrite(os.path.join(args.output_folder,
args.prefix,
"right_image{:05d}.png".format(right_image_iter)),
image)
if right_image_iter > 0:
right_image_times.append(time)
else:
first_right_image_time = time
right_event_image_times.append(time.to_sec())
right_image_iter += 1
elif 'events' in topic and msg.events:
for event in msg.events:
ts = event.ts
event = [event.x,
event.y,
(ts - t_start_ros).to_sec(),
(float(event.polarity) - 0.5) * 2]
if isLeft:
if first_left_image_time != -1 and ts > first_left_image_time:
left_events.append(event)
else:
left_start_event_offset += 1
else:
if first_right_image_time != -1 and ts > first_right_image_time:
right_events.append(event)
else:
right_start_event_offset += 1
if isLeft:
if len(left_image_times) >= args.max_aug and \
left_events[-1][2] > (left_image_times[args.max_aug - 1] - t_start_ros).to_sec():
left_event_image_iter, consumed = _save_events(left_events,
left_image_times,
left_event_count_images,
left_event_time_images,
left_event_image_times,
left_start_event_offset,
left_image_event_start_idx,
left_image_event_end_idx,
rows,
cols,
args.max_aug,
args.n_skip,
left_event_image_iter,
args.prefix,
'left',
left_tf_writer,
t_start_ros,
h5_left,
whitelist_imageids)
left_start_event_offset += consumed
else:
if len(right_image_times) >= args.max_aug and \
right_events[-1][2] > (right_image_times[args.max_aug - 1] - t_start_ros).to_sec():
right_event_image_iter, consumed = _save_events(right_events,
right_image_times,
right_event_count_images,
right_event_time_images,
right_event_image_times,
right_start_event_offset,
right_image_event_start_idx,
right_image_event_end_idx,
rows,
cols,
args.max_aug,
args.n_skip,
right_event_image_iter,
args.prefix,
'right',
right_tf_writer,
t_start_ros,
h5_right,
whitelist_imageids)
right_start_event_offset += consumed
left_tf_writer.close()
right_tf_writer.close()
image_counter_file = open(os.path.join(args.output_folder, args.prefix, "n_images.txt"), 'w')
image_counter_file.write("{} {}".format(left_event_image_iter, right_event_image_iter))
image_counter_file.close()
def main():
parser = argparse.ArgumentParser(
description=("Extracts grayscale and event images from a ROS bag and "
"saves them as TFRecords for training in TensorFlow."))
parser.add_argument("--bag",
dest="bag",
help="Path to ROS bag.",
required=True)
parser.add_argument("--prefix",
dest="prefix",
help="Output file prefix.",
required=True)
parser.add_argument("--output_folder",
dest="output_folder",
help="Output folder.",
required=True)
parser.add_argument(
"--max_aug",
dest="max_aug",
help="Maximum number of images to combine for augmentation.",
type=int,
default=6)
parser.add_argument(
"--n_skip",
dest="n_skip",
help="Maximum number of images to combine for augmentation.",
type=int,
default=1)
parser.add_argument("--start_time",
dest="start_time",
help="Time to start in the bag.",
type=float,
default=0.0)
parser.add_argument("--end_time",
dest="end_time",
help="Time to end in the bag.",
type=float,
default=-1.0)
args = parser.parse_args()
bridge = CvBridge()
n_msgs = 0
left_event_image_iter = 0
right_event_image_iter = 0
left_image_iter = 0
right_image_iter = 0
first_left_image_time = -1
first_right_image_time = -1
left_events = []
right_events = []
left_images = []
right_images = []
left_image_times = []
right_image_times = []
left_event_count_images = []
left_event_time_images = []
left_event_image_times = []
right_event_count_images = []
right_event_time_images = []
right_event_image_times = []
cols = 346
rows = 260
print("Processing bag")
bag = Bag(args.bag)
left_tf_writer = tf.python_io.TFRecordWriter(
os.path.join(args.output_folder, args.prefix,
"left_event_images.tfrecord"))
right_tf_writer = tf.python_io.TFRecordWriter(
os.path.join(args.output_folder, args.prefix,
"right_event_images.tfrecord"))
# Get actual time for the start of the bag.
t_start = bag.get_start_time()
t_start_ros = rospy.Time(t_start)
# Set the time at which the bag reading should end.
if args.end_time == -1.0:
t_end = bag.get_end_time()
else:
t_end = t_start + args.end_time
eps = 0.1
for topic, msg, t in bag.read_messages(
topics=[
'/davis/left/image_raw', '/davis/right/image_raw',
'/davis/left/events', '/davis/right/events'
],
start_time=rospy.Time(max(args.start_time, eps) - eps + t_start),
end_time=rospy.Time(t_end)):
# Check to make sure we're working with stereo messages.
if not ('left' in topic or 'right' in topic):
print(
'ERROR: topic {} does not contain left or right, is this stereo?'
'If not, you will need to modify the topic names in the code.'.
format(topic))
return
# Counter for status updates.
n_msgs += 1
if n_msgs % 500 == 0:
print("Processed {} msgs, {} images, time is {}.".format(
n_msgs, left_event_image_iter,
t.to_sec() - t_start))
isLeft = 'left' in topic
if 'image' in topic:
width = msg.width
height = msg.height
if width != cols or height != rows:
print(
"Image dimensions are not what we expected: set: ({} {}) vs got:({} {})"
.format(cols, rows, width, height))
return
time = msg.header.stamp
if time.to_sec() - t_start < args.start_time:
continue
image = np.asarray(bridge.imgmsg_to_cv2(msg, msg.encoding))
image = np.reshape(image, (height, width))
if isLeft:
cv2.imwrite(
os.path.join(
args.output_folder, args.prefix,
"left_image{:05d}.png".format(left_image_iter)), image)
if left_image_iter > 0:
left_image_times.append(time)
else:
first_left_image_time = time
left_event_image_times.append(time.to_sec())
# filter events we added previously
left_events = filter_events(
left_events, left_event_image_times[-1] - t_start)
left_image_iter += 1
else:
cv2.imwrite(
os.path.join(
args.output_folder, args.prefix,
"right_image{:05d}.png".format(right_image_iter)),
image)
if right_image_iter > 0:
right_image_times.append(time)
else:
first_right_image_time = time
right_event_image_times.append(time.to_sec())
# filter events we added previously
right_events = filter_events(
right_events, right_event_image_times[-1] - t_start)
right_image_iter += 1
elif 'events' in topic and msg.events:
# Add events to list.
for event in msg.events:
ts = event.ts
event = [
event.x, event.y, (ts - t_start_ros).to_sec(),
(float(event.polarity) - 0.5) * 2
]
if isLeft:
# add event if it was after the first image or we haven't seen the first image
if first_left_image_time == -1 or ts > first_left_image_time:
left_events.append(event)
elif first_right_image_time == -1 or ts > first_right_image_time:
right_events.append(event)
if isLeft:
if len(left_image_times) >= args.max_aug and\
left_events[-1][2] > (left_image_times[args.max_aug-1]-t_start_ros).to_sec():
left_event_image_iter = _save_events(
left_events, left_image_times, left_event_count_images,
left_event_time_images, left_event_image_times, rows,
cols, args.max_aug, args.n_skip, left_event_image_iter,
args.prefix, 'left', left_tf_writer, t_start_ros)
else:
if len(right_image_times) >= args.max_aug and\
right_events[-1][2] > (right_image_times[args.max_aug-1]-t_start_ros).to_sec():
right_event_image_iter = _save_events(
right_events, right_image_times,
right_event_count_images, right_event_time_images,
right_event_image_times, rows, cols, args.max_aug,
args.n_skip, right_event_image_iter, args.prefix,
'right', right_tf_writer, t_start_ros)
left_tf_writer.close()
right_tf_writer.close()
image_counter_file = open(
os.path.join(args.output_folder, args.prefix, "n_images.txt"), 'w')
image_counter_file.write("{} {}".format(left_event_image_iter,
right_event_image_iter))
image_counter_file.close()
