def crop_bag(bag_file, annotation_file):
if not os.path.exists(annotation_file):
print 'no annotations - skipping'
return
annotation_bag = rosbag.Bag(annotation_file, 'r')
bag = rosbag.Bag(bag_file.replace('compressed_',''), 'r')
wheelchair_bag_info = get_info(bag)
wheelchair_topics = [i['topic'] for i in wheelchair_bag_info['topics']]
annotation_bag_info = get_info(annotation_bag)
annotation_topics = [i['topic'] for i in annotation_bag_info['topics']]
if '/bag_interval' in annotation_topics:
annotation_topics.remove('/bag_interval')
intervals = []
for _, msg, time in annotation_bag.read_messages(topics=['/bag_interval']):
intervals.append(dur_to_startend(time, msg.data))
out_bag_path = annotation_path(bag_file, 'cropped')
out_bag = rosbag.Bag(out_bag_path, 'w')
extract_clusters(intervals, wheelchair_topics, bag, out_bag)
if annotation_topics:
for data in annotation_bag.read_messages(topics=annotation_topics):
out_bag.write(*data)
out_bag.close()
print 'wrote', out_bag_path
ind_r += 1
elif topic.endswith("wifi_info"):
last_wifi_status = len(msg.essid) > 0
bags[os.path.split(bagfile)[-1]] = dict(left=ssd(np.ma.asarray(left)), right=ssd(np.ma.asarray(right)))
ssd_left = bags[os.path.split(bagfile)[-1]]["left"]
ssd_right = bags[os.path.split(bagfile)[-1]]["right"]
l_indices_to_keep = (ssd_left > 35) & wifi_left[:-1]
r_indices_to_keep = (ssd_right > 35) & wifi_right[:-1]
l_stamps = np.array(l_stamps)[l_indices_to_keep.nonzero()[0]]
r_stamps = np.array(r_stamps)[r_indices_to_keep.nonzero()[0]]
if l_indices_to_keep.sum() + r_indices_to_keep.sum() > 0:
clusters = find_clusters(np.concatenate((l_stamps, r_stamps)), dist=0.5)
left_keep_scans = left[l_indices_to_keep]
right_keep_scans = right[r_indices_to_keep]
path, filename = os.path.split(bagfile)
bag_out = rosbag.Bag(os.path.join(path, "compressed_" + filename), "w", compression=rosbag.Compression.BZ2)
extract_clusters(clusters, topics, bag, bag_out)
bag_out.close()
print "wrote ", bag_out._filename
else:
print "no movement - skipping bag"
