def gen_TFRecord_from_file(out_dir, out_filename, bag_filename, flip=False):
packager = DataPackager(flip=flip)
bag = rosbag.Bag(bag_filename)
output_filenames = []
#######################
## Get Label Info ##
#######################
example_id = out_filename
file_end = bag_filename.find(".bag")
label_code = bag_filename[file_end-5:file_end]
print("")
print("bag_filename: ", bag_filename)
print("label_code:", label_code)
img_lab, opt_lab, aud_lab = 0,0,0
if("z" in example_id):
img_lab = 1
if("g" in example_id):
opt_lab = 1
if("a" in example_id):
aud_lab = 1
total_lab = (img_lab+opt_lab+aud_lab > 0)
print(example_id)
print(img_lab, opt_lab, aud_lab, ':', total_lab)
end_file = ".tfrecord"
if(flip):
end_file = "_flip"+end_file
#######################
## READ FILE ##
#######################
p_t = 0
stored_data = []
for topic, msg, t in bag.read_messages(topics=topic_names):
if(topic == topic_names[0]):
last_action = str(msg.data)
if(msg.data > 0):
# perform data pre-processing steps
packager.formatOutput()
if(msg.data == 1):
print("packager.getImgStack().shape: ", packager.getImgStack().shape)
stored_data = {
"img_raw": packager.getImgStack()[:], "img_lab": 0,
"aud_raw": packager.getAudStack()[:], "aud_lab": 0,
"p_t": p_t,
"total_lab": int(last_action),
"example_id": example_id}
p_t += 1
elif(msg.data > 1):
break
packager.reset()
elif(topic == topic_names[1]):
packager.imgCallback(msg)
elif(topic == topic_names[2]):
packager.audCallback(msg)
if(p_t > 0):
ex = make_sequence_example (
img_raw=stored_data["img_raw"], img_lab=stored_data["img_lab"],
aud_raw=stored_data["aud_raw"], aud_lab=stored_data["aud_lab"],
p_t=stored_data["p_t"],
first_action=stored_data["total_lab"],
example_id=stored_data["example_id"],
img_raw2=packager.getImgStack(),
aud_raw2=packager.getAudStack(),
second_action=int(last_action))
output_filename = out_dir+out_filename+"_"+str(stored_data["total_lab"])+end_file
output_filenames.append(output_filename)
writer = tf.python_io.TFRecordWriter(output_filename)
writer.write(ex.SerializeToString())
writer.close()
# generate TFRecord data
ex = make_sequence_example (
img_raw=packager.getImgStack(), img_lab=img_lab,
aud_raw=packager.getAudStack(), aud_lab=aud_lab,
p_t=p_t,
first_action=int(last_action),
example_id=example_id)
print("last_action:", msg.data, int(last_action))
# write TFRecord data to file
output_filename = out_dir+out_filename+"_"+last_action+end_file
output_filenames.append(output_filename)
writer = tf.python_io.TFRecordWriter(output_filename)
writer.write(ex.SerializeToString())
writer.close()
packager.reset()
bag.close()
return output_filenames
