def preprocess(bagfile, output, output_size): """ yields topic, data, t but data already a numpy array """ bag = rosbag.Bag(bagfile) info = yaml.load(bag._get_yaml_info()) for topic_dict in info['topics']: if 'FloatArray' in topic_dict['type']: command_topic = topic_dict['topic'] topics = [topic_image_raw, command_topic] logger.info('bag_processing, using topics: %s' % topics) def read_tuples(raw_stream): i = 0 Y0 = None U0 = None Y_last = None t0 = 0 for topic, msg, t in raw_stream: # logger.debug('read msg type %s and topic %s' % (msg._type, topic)) i += 1 if topic == command_topic: # logger.debug('command read') Y1 = Y_last if Y0 is not None and U0 is not None: # yeild tuple yield (Y0, U0, Y1, t0) Y0 = Y_last U0 = msg t0 = t if topic == topic_image_raw: # logger.debug('image read') Y_last = msg tuples_stream = read_tuples(bag.read_messages(topics=topics)) out_bag = rosbag.Bag(output, 'w') for Y0, U, Y1, t0 in tuples_stream: y0_pil = Image.fromarray(get_image_array(Y0)).resize(output_size) y1_pil = Image.fromarray(get_image_array(Y1)).resize(output_size) # pdb.set_trace() logger.info('Writing data (Y0, U0=%s, Y1)' % U) out_bag.write('Y0', pil_to_imgmsg(y0_pil), t0) out_bag.write('U0', U, t0) out_bag.write('Y1', pil_to_imgmsg(y1_pil), t0) out_bag.close()
def write_tuple_to_bag(bag, Y0, U, Y1, x=0): # write to the bag y0msg = pil_to_imgmsg(Image.fromarray(Y0[1])) y0msg.header.stamp = U[0] y1msg = pil_to_imgmsg(Image.fromarray(Y1[1])) y1msg.header.stamp = U[0] # pdb.set_trace() bag.write('Y0', y0msg, U[0]) bag.write('U0', U[1], U[0]) bag.write('Y1', y1msg, U[0]) try: bag.write('X0', std_msgs.msg.Int64(int(x)), U[0]) except: pdb.set_trace()