def process_bag(bag_in_fn, bag_out_fn, conf_file_fn): bag_in = Bag(bag_in_fn) bag_out = Bag(bag_out_fn, 'w') include_rules, exclude_rules, time_rules = read_rules(conf_file_fn) topic_rules = include_rules + exclude_rules # Set the time UNIX time at the start of the bag file for r in topic_rules: r.set_begin_time(bag_in.get_start_time()) for r in time_rules: r.set_begin_time(bag_in.get_start_time()) # Find start and end times that are actually required t_start, t_end = get_begin_end(time_rules, bag_in) # Find topics that are actually required bag_topics = bag_in.get_type_and_topic_info().topics.keys() topics = get_topics(topic_rules, bag_topics) messages = bag_in.read_messages(topics=topics, start_time=t_start, end_time=t_end, return_connection_header=True) for topic, msg, t, conn_header in messages: # Check default enforcement for this message if FilterRule.is_tf_topic(topic): default = FilterRule.DEFAULT_ENFORCEMENT_TF else: default = FilterRule.DEFAULT_ENFORCEMENT_TOPIC # When default is to include, only check whether the exclusion # rules are satisfied, and if all of them are ok, write it out if default == FilterRule.INCLUDE: # Check exclusions ok = True for r in exclude_rules: if r.msg_match(topic, msg, t): ok = False # When default is to exclude, check if the message matches any # of the inclusion rules and write it out if it does else: # Check inclusions ok = False for r in include_rules: if r.msg_match(topic, msg, t): ok = True # Time rules time_ok = True for r in time_rules: if not r.is_ok_with(topic, msg, t): time_ok = False # Write to file if ok and time_ok: bag_out.write(topic, msg, t, connection_header=conn_header) bag_out.close()
def main(): parser = argparse.ArgumentParser() parser.add_argument("--bag", dest="bag", help="Path to ROS bag.", required=True) parser.add_argument("--left_image_0", help="Path to 'left_image_00000.png'") parser.add_argument("--right_image_0", help="Path to 'right_image_00000.png'") parser.add_argument("--start_time", help="A start time to check", type=float, default=None) args = parser.parse_args() bridge = CvBridge() bag = Bag(args.bag) t_start = bag.get_start_time() n_msgs = 0 if args.left_image_0 is None and args.right_image_0 is None: raise ValueError( "You must provide either left_image_0 or right_image_0") isLeft = args.left_image_0 is not None ref_img = args.left_image_0 if isLeft else args.right_image_0 print("Using ref image %s" % ref_img) ref_image_0 = cv2.imread(ref_img)[:, :, 0] assert ref_image_0 is not None topic = '/davis/left/image_raw' if isLeft else '/davis/right/image_raw' if args.start_time: topic, msg, t = next( bag.read_messages(topics=[topic], start_time=rospy.Time(args.start_time + t_start))) width = msg.width height = msg.height image = np.asarray(bridge.imgmsg_to_cv2(msg, msg.encoding)) image = np.reshape(image, (height, width)) print("Correct? {}".format(bool(np.all(np.isclose(image, ref_image_0))))) else: for topic, msg, t in bag.read_messages(topics=[topic]): n_msgs += 1 width = msg.width height = msg.height image = np.asarray(bridge.imgmsg_to_cv2(msg, msg.encoding)) image = np.reshape(image, (height, width)) if bool(np.all(np.isclose(image, ref_image_0))): print("Ref image found at {} ({} from start)".format( t.to_sec(), t.to_sec() - t_start)) return print("Processed {} images, ref not found!".format(n_msgs))
def bag_to_video( input_file: Bag, output_file: str, topic: str, fps: float = 30, codec: str = 'MJPG', ) -> None: """ Convert a ROS bag with image topic to a video file. Args: input_file: the bag file to get image data from output_file: the path to an output video file to create topic: the topic to read image data from fps: the frame rate of the video codec: the codec to use when outputting to the video file Returns: None """ # create an empty reference for the output video file video = None # get the total number of frames to write total = input_file.get_message_count(topic_filters=topic) # get an iterator for the topic with the frame data iterator = input_file.read_messages(topics=topic) # iterate over the image messages of the given topic for _, msg, _ in tqdm(iterator, total=total): # open the video file if it isn't open if video is None: # create the video codec codec = cv2.VideoWriter_fourcc(*codec) # open the output video file cv_dims = (msg.width, msg.height) video = cv2.VideoWriter(output_file, codec, fps, cv_dims) # read the image data into a NumPy tensor img = get_camera_image(msg.data, (msg.height, msg.width)) # write the image to the video file video.write(img) # if the video file is open, close it if video is not None: video.release() # read the start time of the video from the bag file start = input_file.get_start_time() # set the access and modification times for the video file os.utime(output_file, (start, start))
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()
def semantic_segment( metadata: str, input_bag: Bag, model: 'keras.models.Model', predict: str, output_bag: Bag = None, output_dir: str = None, base: str = None, num_samples: int = 200, encoding: str = 'rgb', ) -> None: """ Predict a stream of images from an input ROSbag. Args: metadata: the metadata about the semantic segmentations from the model input_bag: the input bag to predict targets from a topic model: the semantic segmentation model to use to make predictions predict: the topic to get a priori estimates from output_bag: the output bag to write the a priori estimates to output_dir: the output directory to write image pairs to base: the base-name for the prediction image topic num_samples: the number of image pairs to sample for output directory encoding: the encoding for the images to write Returns: None """ # create the base endpoint for the topics base = '' if base is None else '{}'.format(base) # setup the output directories if output_dir is not None: x_dir = os.path.join(output_dir, 'X', 'data') if not os.path.isdir(x_dir): os.makedirs(x_dir) y_dir = os.path.join(output_dir, 'y', 'data') if not os.path.isdir(y_dir): os.makedirs(y_dir) # read the RGB map and vectorized method from the metadata file rgb_map, unmap_rgb = read_rgb_map(metadata) # write the color map metadata to the output bag if output_bag is not None: ros_stamp = rospy.rostime.Time(input_bag.get_start_time()) msg = String(repr(rgb_map)) output_bag.write('{}/rgb_map'.format(rgb_map), msg, ros_stamp) # open a Window to play the video x_window = Window('img', model.input_shape[1], model.input_shape[2]) y_window = Window('sem-seg', model.output_shape[1], model.output_shape[2]) # create a progress bar for iterating over the messages in the bag total_messages = input_bag.get_message_count(topic_filters=predict) with tqdm(total=total_messages, unit='message') as prog: # iterate over the messages in this input bag for _, msg, time in input_bag.read_messages(topics=predict): # update the progress bar with a single iteration prog.update(1) if np.random.random() > num_samples / total_messages: continue # create a tensor from the raw pixel data pixels = get_camera_image(msg.data, (msg.height, msg.width))[..., :3] # flip the BGR image to RGB if encoding == 'bgr': pixels = pixels[..., ::-1] # resize the pixels to the shape of the model _pixels = resize( pixels, model.input_shape[1:], anti_aliasing=False, mode='symmetric', clip=False, preserve_range=True, ).astype('uint8') # pass the frame through the model y_pred = model.predict(_pixels[None, ...])[0] y_pred = np.stack(unmap_rgb(y_pred.argmax(axis=-1)), axis=-1) y_pred = y_pred.astype('uint8') # show the pixels on the windows x_window.show(_pixels) y_window.show(y_pred) # create an Image message and write it to the output ROSbag if output_bag is not None: msg = image_msg(y_pred, msg.header.stamp, y_pred.shape[:2], 'rgb8') output_bag.write('{}/image_raw'.format(base), msg, msg.header.stamp) # sample a number and write the image pair to disk if output_dir is not None: x_file = os.path.join(x_dir, '{}.png'.format(time)) Image.fromarray(pixels).save(x_file) y_file = os.path.join(y_dir, '{}.png'.format(time)) y_pred = resize( y_pred, pixels.shape[:2], anti_aliasing=False, mode='symmetric', clip=False, preserve_range=True, ).astype('uint8') Image.fromarray(y_pred).save(y_file)