def bag_to_clock(filename='clock.bag'):
bagfile = Bag(filename)
clock = []
for topic, msg, timestamp in bagfile.read_messages():
if topic == '/clock':
clock.append([timestamp.to_sec(), msg.clock.to_sec()])
return numpy.array(clock)
def bag_to_power(filename='power.bag'):
bagfile = Bag(filename)
energies = {}
power = []
for topic, msg, timestamp in bagfile.read_messages():
if topic.find('energy') != -1:
energies[topic] = msg.data
power.append([timestamp.to_sec(), sum(energies.values())])
bagfile.close()
return numpy.array(power)
def bag_to_csv(filename):
bagfile = Bag(filename)
topics = init_topics(filename)
start_time = bagfile.get_start_time()
for topic, msg, timestamp in bagfile.read_messages(topics=topics.keys()):
formatter = topics[topic]['formatter']
data_line = formatter((topic,msg,timestamp,start_time))
topics[topic]['file'].write(', '.join(map(str,data_line)) + '\n')
for topic in topics.keys():
topics[topic]['file'].close()
bagfile.close()
def bag_to_blobs(filename):
bagfile = Bag(filename)
blobs = []
b_time = None
curr_blobs = [None, None]
poses = []
start_time = bagfile.get_start_time()
for topic, msg, timestamp in bagfile.read_messages():
if topic[-9:] == 'blob_list':
if len(msg.blobs) == 1:
m_time = msg.header.stamp.to_sec()
r_idx = int(topic[-20])
blob = [msg.blobs[0].x, msg.blobs[0].y]
#blobs.append([m_time,r_idx] + blob)
if b_time is None:
b_time = m_time
curr_blobs[r_idx] = blob
elif abs(b_time - m_time) < 0.001:
curr_blobs[r_idx] = blob
blobs.append([b_time] + curr_blobs[0] + curr_blobs[1])
b_time = None
curr_blobs = [None, None]
else:
b_time = m_time
curr_blobs = [None, None]
curr_blobs[r_idx] = blob
elif topic == '/tf':
for tf_msg in msg.transforms:
p_time = tf_msg.header.stamp.to_sec()
base_frame = tf_msg.header.frame_id
child_frame = tf_msg.child_frame_id
pos = tf_msg.transform.translation
pos = [pos.x, pos.y, pos.z]
ori = tf_msg.transform.rotation
ori = [ori.x, ori.y, ori.z, ori.w]
poses.append([p_time, base_frame, child_frame, pos, ori])
bagfile.close()
return blobs, poses
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="save data from bag file")
parser.add_argument("--bag", required=True, help="Path to bag file")
parser.add_argument("--output", default="bag_output", help="directory for data")
parser.add_argument("--blob-topic", dest="blob_topic",
default="/gatlin/objectscreencoords",
help="topic for blobs")
parser.add_argument("--image-topic", dest="image_topic",
default="/camera/rgb/image_rect_color_throttled",
help="topic for images")
parser.add_argument("--factor", default=10**7, help="Factor for decimal digits")
parser.add_argument("--interactive", action="store_true", help="Show all images to mark blocks")
args = parser.parse_args()
bag = Bag(args.bag)
messages = bag.read_messages()
# Convert bag data to correct format
groups = defaultdict(lambda : defaultdict(lambda : defaultdict(list)))
blobs_to_group_later = []
for topic, msg, time in messages:
if topic == args.blob_topic:
data = [msg.x, msg.y]
blobs_to_group_later.append((data, time))
else:
data = np.fromstring(msg.data, np.uint8)
data = np.reshape(data, (msg.height, msg.width, 3))
groups[time.secs][int(time.nsecs / args.factor)][args.image_topic].append(data)
# Match blobs with images
for data, time in blobs_to_group_later:
class RosbagReader():
def __init__(self, input_bag, save_dir):
self.save_dir = save_dir
Path.mkdir(self.save_dir, parents=True, exist_ok=True)
self.bag = Bag(input_bag)
def print_bag_info(self):
info_dict = yaml.load(Bag(self.bag, 'r')._get_yaml_info())
print(info_dict)
def extract_camera_data(self):
image_topic = '/camera_left/color/image_raw'
bag_name = self.bag.filename.strip(".bag").split('/')[-1]
output_dir = os.path.join(self.save_dir, bag_name, 'camera')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
bridge = CvBridge()
count = 0
for topic, msg, t in self.bag.read_messages(topics=[image_topic]):
cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
cv2.imwrite(os.path.join(output_dir, "frame%06i.png" % count),
cv_img)
print("Wrote image {}".format(count))
count += 1
def extract_lidar_frames(self, extension, topicList=None):
topic_name = '/front_lidar/velodyne_points'
# bag_name = self.bag.filename.strip(".bag").split('/')[-1]
# save_temp_dir = os.path.join(self.save_dir,bag_name)
i = 0
for topic, msg, t in self.bag.read_messages(topic_name):
# for each instance in time that has data for topicName
# parse data from this instance
save_file = self.save_dir / 'frame_{}.{}'.format(i, extension)
i += 1
if topic == '/front_lidar/velodyne_points':
pc_list = [[p[0], p[1], p[2], p[3]]
for p in pc2.read_points(msg,
skip_nans=True,
field_names=("x", "y", "z",
"intensity"))]
# print(np.array(pc_list, dtype=np.float32).shape)
def read_bag(self, save_to_csv=True):
"""
Return dict with all recorded messages with topics as keys
Args:
save_to_csv (bool, optional): Save data to csv files (one individual file per topic) if True
Returns:
dict: containing all published data points per topic
"""
topics = self.readBagTopicList()
messages = {}
max_it_bag = 10
it_bag = 0
# iterate through topic, message, timestamp of published messages in bag
for topic, msg, _ in self.bag.read_messages(topics=topics):
if type(msg).__name__ == '_sensor_msgs__PointCloud2':
points = np.zeros((msg.height * msg.width, 3))
for pp, ii in zip(
pc2.read_points(msg,
skip_nans=True,
field_names=("x", "y", "z")),
range(points.shape[0])):
points[ii, :] = [pp[0], pp[1], pp[2]]
msg = points
if topic not in messages:
messages[topic] = [msg]
else:
messages[topic].append(msg)
it_bag += 1
if it_bag >= max_it_bag:
break # stop -- too many topics
return messages
def readBagTopicList(self):
"""
Read and save the initial topic list from bag
"""
print("[OK] Reading topics in this bag. Can take a while..")
topicList = []
bag_info = yaml.load(self.bag._get_yaml_info(), Loader=yaml.FullLoader)
for info in bag_info['topics']:
topicList.append(info['topic'])
print('{0} topics found'.format(len(topicList)))
return topicList
class RosbagReader():
def __init__(self, input_bag, save_dir):
self.save_dir = save_dir
Path.mkdir(self.save_dir, parents=True, exist_ok=True)
self.bag = Bag(input_bag)
def print_bag_info(self):
info_dict = yaml.load(Bag(self.bag, 'r')._get_yaml_info())
print(info_dict)
def extract_camera_data(self):
image_topic = '/camera_left/color/image_raw'
bag_name = self.bag.filename.strip(".bag").split('/')[-1]
output_dir = os.path.join(self.save_dir, bag_name, 'camera')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
bridge = CvBridge()
count = 0
for topic, msg, t in self.bag.read_messages(topics=[image_topic]):
cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
cv2.imwrite(os.path.join(output_dir, "frame%06i.png" % count),
cv_img)
print("Wrote image {}".format(count))
count += 1
def extract_lidar_frames(self, extension, topicList=None):
topic_name = '/front_lidar/velodyne_points'
# bag_name = self.bag.filename.strip(".bag").split('/')[-1]
# save_temp_dir = os.path.join(self.save_dir,bag_name)
i = 0
for topic, msg, t in self.bag.read_messages(topic_name):
# for each instance in time that has data for topicName
# parse data from this instance
save_file = self.save_dir / 'frame_{}.{}'.format(i, extension)
i += 1
if topic == '/front_lidar/velodyne_points':
pc_list = [[p[0], p[1], p[2], p[3]]
for p in pc2.read_points(msg,
skip_nans=True,
field_names=("x", "y", "z",
"intensity"))]
# print(np.array(pc_list, dtype=np.float32).shape)
if extension == 'pcd':
p = pcl.PointCloud(np.array(pc_list, dtype=np.float32))
p.to_file(save_file)
elif extension == 'h5':
save_h5_basic(save_file, np.array(pc_list, dtype=np.float32))
print('---> Saved {} frames to files.'.format(i))
def save_bag_to_csv(self, topicList=None):
if not topicList:
topicList = self.readBagTopicList()
for topic_name in topicList:
bag_name = self.bag.filename.strip(".bag").split('/')[-1]
filename = self.save_dir + bag_name + topic_name.replace(
'/', '-') + '.csv'
# create csv file for topic data
with open(filename, 'w+') as csvfile:
print("Parsing data. This may take a while.")
filewriter = csv.writer(csvfile, delimiter=',')
flag = 1 # use flag to only write header once to csv file
for topic, msg, t in self.bag.read_messages(topic_name):
# for each instance in time that has data for topicName
# parse data from this instance
row = [t]
if topic == '/front_lidar/velodyne_points':
header = [
"rosbagTimestamp", "x", "y", "z", "intensity",
"time"
]
if flag:
filewriter.writerow(header)
flag = 0
for p in pc2.read_points(msg,
skip_nans=True,
field_names=("x", "y", "z",
"intensity",
"time")):
row += [p[0], p[1], p[2], p[3], p[4]]
filewriter.writerow(row)
row = [t]
else:
msg_fields = msg.fields
field_names = []
header = ['rosbagTimestamp']
for f in msg_fields:
field_names.append(f.name)
header.append(f.name)
if flag:
filewriter.writerow(header)
flag = 0
for p in pc2.read_points(
msg,
skip_nans=True,
field_names=tuple(field_names)):
row += list(p)
filewriter.writerow(row)
row = [t]
print("Parsed data. Saved as {}".format(filename))
def save_rwth_detections(self):
"""
Saves Detections generated by RWTH Detector into separate csv files.
"""
os.chdir(self.save_dir)
topic_name = '/drow_detected_persons'
bag_name = string.rstrip(self.bag.filename, ".bag")
filename = self.save_dir + bag_name + string.replace(
topic_name, '/', '-') + '.csv'
# create csv file for topic data
with open(filename, 'w+') as csvfile:
print("Parsing data. This may take a while.")
filewriter = csv.writer(csvfile, delimiter=',')
header = [
'rosbagTimestamp', 'detection_id', 'confidence', 'det_x',
'det_y', 'det_z', 'height'
]
filewriter.writerow(header)
for _, msg, t in self.bag.read_messages(topic_name):
for det in msg.detections:
row = [t]
row.append(det.detection_id)
row.append(det.confidence)
row.append(det.pose.pose.position.x)
row.append(det.pose.pose.position.y)
row.append(det.pose.pose.position.z)
row.append(det.height)
filewriter.writerow(row)
row = [t]
print("Parsed data. Saved as {}".format(filename))
def read_bag(self, save_to_csv=True):
"""
Return dict with all recorded messages with topics as keys
Args:
save_to_csv (bool, optional): Save data to csv files (one individual file per topic) if True
Returns:
dict: containing all published data points per topic
"""
topics = self.readBagTopicList()
messages = {}
max_it_bag = 10
it_bag = 0
# iterate through topic, message, timestamp of published messages in bag
for topic, msg, _ in self.bag.read_messages(topics=topics):
if type(msg).__name__ == '_sensor_msgs__PointCloud2':
points = np.zeros((msg.height * msg.width, 3))
for pp, ii in zip(
pc2.read_points(msg,
skip_nans=True,
field_names=("x", "y", "z")),
range(points.shape[0])):
points[ii, :] = [pp[0], pp[1], pp[2]]
msg = points
if topic not in messages:
messages[topic] = [msg]
else:
messages[topic].append(msg)
it_bag += 1
if it_bag >= max_it_bag:
break # stop -- too many topics
return messages
def readBagTopicList(self):
"""
Read and save the initial topic list from bag
"""
print("[OK] Reading topics in this bag. Can take a while..")
topicList = []
bag_info = yaml.load(self.bag._get_yaml_info(), Loader=yaml.FullLoader)
for info in bag_info['topics']:
topicList.append(info['topic'])
print('{0} topics found'.format(len(topicList)))
return topicList