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()