def load_trajectory_from_bag(file,fast=True,fast_only=False): traj = [] fast_traj = [] bag = Bag(file) normal_topics = ['raven_state', 'raven_state/array'] normal_topics = normal_topics + ['/' + topic for topic in normal_topics] fast_topic = ['raven_state/1000Hz'] #fast_topic = fast_topic + ['/' + topic for topic in fast_topic] if fast_only: for topic, msg, t in bag.read_messages(topics=fast_topic): fast_traj.append(msg) return get_trajectory(fast_traj) for topic, msg, t in bag.read_messages(topics=normal_topics): traj.append(msg) if fast: for topic, msg, t in bag.read_messages(topics=fast_topic): fast_traj.append(msg) if fast_traj: return get_trajectory(fast_traj,traj) else: return get_trajectory(traj)
def load_trajectory_from_bag(file, fast=True, fast_only=False):
traj = []
fast_traj = []
bag = Bag(file)
normal_topics = ['raven_state', 'raven_state/array']
normal_topics = normal_topics + ['/' + topic for topic in normal_topics]
fast_topic = ['raven_state/1000Hz']
#fast_topic = fast_topic + ['/' + topic for topic in fast_topic]
if fast_only:
for topic, msg, t in bag.read_messages(topics=fast_topic):
fast_traj.append(msg)
return get_trajectory(fast_traj)
for topic, msg, t in bag.read_messages(topics=normal_topics):
traj.append(msg)
if fast:
for topic, msg, t in bag.read_messages(topics=fast_topic):
fast_traj.append(msg)
if fast_traj:
return get_trajectory(fast_traj, traj)
else:
return get_trajectory(traj)
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 get_start_stamp(bag: rosbag.Bag):
pos_prev = []
t_prev = 0
count = 0
t_start = 0
for topic, msg, t in bag.read_messages('/tf'):
if msg.transforms[0].child_frame_id == "Puck":
pos = np.array([
msg.transforms[0].transform.translation.x,
msg.transforms[0].transform.translation.y,
msg.transforms[0].transform.translation.z
])
t_sec = msg.transforms[0].header.stamp.to_sec()
if t_prev != 0:
vel = np.linalg.norm((pos - pos_prev) / (t_sec - t_prev))
if vel > 0.1:
if count == 0:
t_start = t
count += 1
if count > 10:
return t_start
else:
count = 0
pos_prev = pos
t_prev = t_sec
def get_info(bag_file, topic_filter=None):
bag = Bag(bag_file)
topics = bag.get_type_and_topic_info().topics
for topic in topics:
if topic_filter and topics[topic].msg_type not in topic_filter:
continue
print("{}: {} Hz".format(topic, round(topics[topic].frequency, 3)))
print(topics[topic].message_count)
times = np.ones(shape=bag.get_message_count(topic_filters=topic))
n = 0
for _, msg, t in bag.read_messages(topics=topic):
times[n] = msg.header.stamp.to_sec()
n += 1
times = 1 / np.gradient(times)
times = times[np.where((times > np.percentile(times, 10)) & (times < np.percentile(times, 90)))]
print("mean: {}, median: {}".format(np.mean(times), np.median(times)))
print("min: {}, max: {}".format(np.min(times), np.max(times)))
# plt.scatter(times, np.gradient(times))
plt.hist(times)
plt.yscale("log")
plt.title("{}: {}".format(os.path.basename(bag_file), topic))
if not os.path.exists("images"):
os.makedirs("images")
plt.savefig(os.path.join("images/", "{}.{}.png".format(os.path.basename(bag_file), topic.replace("/", ""))))
plt.cla()
def dump(input_file: Bag, topics: list, output_file: 'file' = None) -> None:
"""
Dump messages from a bag.
Args:
input_file: the input bag to dump topics from
topics: the topics to dump
output_file: an optional file to dump to
Returns:
None
"""
# create a progress bar for iterating over the messages in the bag
with tqdm(total=input_file.get_message_count(
topic_filters=topics)) as prog:
# iterate over the messages in this input bag
for topic, msg, _ in input_file.read_messages(topics=topics):
# update the progress bar with a single iteration
prog.update(1)
# create the line to print
line = '{} {}\n\n'.format(topic, msg)
# print the line to the terminal
print(line)
# if there is an output file, write the line to it
if output_file is not None:
output_file.write(line)
def _assert_bag_valid(self,
filename,
topics=None,
start_time=None,
stop_time=None):
'''
Open the bagfile at the specified filename and read it to ensure topic limits were
enforced and the optional topic list and start/stop times are also enforced.
'''
bag = Bag(filename)
topics_dict = bag.get_type_and_topic_info()[1]
bag_topics = set(topics_dict.keys())
param_topics = set(self.topic_limits.keys())
if topics:
self.assertEqual(bag_topics, set(topics))
self.assertTrue(bag_topics.issubset(param_topics))
for topic in topics_dict:
size = topics_dict[
topic].message_count * 8 # Calculate stored message size as each message is 8 bytes
count = topics_dict[topic].message_count
gen = bag.read_messages(topics=topic)
_, _, first_time = next(gen)
last_time = first_time # in case the next for loop does not execute
if start_time:
self.assertGreaterEqual(first_time, start_time)
for _, _, last_time in gen: # Read through all messages so last_time is valid
pass
if stop_time:
self.assertLessEqual(last_time, stop_time)
duration = last_time - first_time
self._assert_limits_enforced(topic, duration, size, count)
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 extract_depth(
bag_file: Bag,
rgb_directory: str,
camera_info: str,
depth: str
) -> None:
"""
Extract depth data from a bag file.
Args:
bag_file: the bag file to play
rgb_directory: the directory to find RGB image to match depths to
camera_info: the topic to use to read metadata about the camera
depth: the topic to use to read 32-bit floating point depth measures
Returns:
None
"""
# extract the camera dimensions from the bag
dims = get_camera_dimensions(bag, camera_info)
# get the images from a glob
images = glob.glob(os.path.join(rgb_directory, 'X', 'data', '*.png'))
# convert the images to numbers
path_to_int = lambda x: int(os.path.basename(x).replace('.png', ''))
images = sorted([path_to_int(image) for image in images])
# create the output directory
output_dir = os.path.join(rgb_directory, 'D', 'data')
try:
os.makedirs(output_dir)
except FileExistsError:
pass
# iterate over the messages
progress = tqdm(total=len(images))
for _, msg, time in bag_file.read_messages(topics=depth):
# if there are no more images left, break out of the loop
if not len(images):
break
# if the time is less than the current image, continue
if int(str(time)) < images[0]:
continue
# update the progress bar
progress.update(1)
# get the depth image
img = get_depth_image(msg.data, dims, as_rgb=False)
# save the depth image to disk
output_file = os.path.join(output_dir, '{}-{}.npz'.format(images[0], time))
np.savez_compressed(output_file, y=img)
# remove the first item from the list of times
images.pop(0)
# close the progress bar
progress.close()
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 ipfs_download(multihash: Multihash) -> (dict, Bag):
rospy.wait_for_service('/ipfs/get_file')
download = rospy.ServiceProxy('/ipfs/get_file', IpfsDownloadFile)
tmpfile = NamedTemporaryFile(delete=False)
res = download(multihash, Filepath(tmpfile.name))
tmpfile.close()
if not res.success:
raise Exception(res.error_msg)
messages = {}
bag = Bag(tmpfile.name, 'r')
for topic, msg, timestamp in bag.read_messages():
if topic not in messages:
messages[topic] = [msg]
else:
messages[topic].append(msg)
os.unlink(tmpfile.name)
return (messages, bag)
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')
topic_rules, tf_rules = RenameRule.parse(conf_file_fn)
messages = bag_in.read_messages(return_connection_header=True)
for topic, msg, t, conn_header in messages:
if topic == '/tf' or topic == '/tf_static':
new_topic = topic
new_msg = modify_tf(msg, tf_rules)
else:
new_topic = modify_topic(topic, topic_rules)
# Modify the frame_id in header if header exists
new_msg = modify_msg(msg, tf_rules)
bag_out.write(new_topic, new_msg, t, connection_header=conn_header)
bag_in.close()
bag_out.close()
def writepc2frombag2file(input_bagfile, pc2_topic, out_file_path,start_time=None, end_time=None, msg_count=None):
try:
output_file_name = get_output_file_name(input_bagfile) + '_xyz.txt'
if out_file_path.endswith('/'):
output_file = out_file_path + output_file_name
else:
output_file = out_file_path+'/' + output_file_name
output_file_fh = open(output_file,'w')
print('writing output_file at', output_file)
# check start and end time condition
if start_time is not None and end_time is not None:
assert (end_time - start_time) > 0 , "end_time should be higher than start time"
if msg_count is None:
use_msg_count = False
else:
use_msg_count = True
# msg_count is already static casted in arg parse, just checking again for sanity ad if function is re used individially
msg_count = int(msg_count)
assert msg_count > 0, "should have positive msg_count"
input_bag = Bag(input_bagfile,'r')
print('bag load success')
count = 0;
for topic, msg, ts in input_bag.read_messages(topics=[pc2_topic],start_time=start_time, end_time=end_time):
output_file_fh.write('msg_timestamp: %f\n' % msg.header.stamp.to_sec())
output_file_fh.write('%d\n' %msg.width)
for data_pts in pc2.read_points(msg, field_names=("x", "y", "z", "intensity"), skip_nans=True):
output_file_fh.write('%f %f %f %d\n' % (data_pts[0], data_pts[1], data_pts[2], data_pts[3]))
count +=1
print(count,msg_count)
if use_msg_count and count >= msg_count:
break
input_bag.close()
output_file_fh.close()
except Exception as e_init:
print(traceback.format_exc(e_init))
def pinata_rosbag(ipfs_hash: str) -> str:
data = requests.get(f'https://gateway.pinata.cloud/ipfs/{ipfs_hash}')
rospy.loginfo(data.status_code)
if data.status_code != 200:
rospy.loginfo(f'Error downloading rosbag! {data.content}')
return
else:
tmpfile = NamedTemporaryFile(delete=False)
tmpfile.write(data.content)
bag = Bag(tmpfile.name)
messages = {}
for topic, msg, timestamp in bag.read_messages():
if topic not in messages:
messages[topic] = [msg]
else:
messages[topic].append(msg)
os.unlink(tmpfile.name)
return (messages, bag)
def get_initial_state(bag: rosbag.Bag):
i = 0
j = 0
tf_table = None
tf_puck_start = None
tf_puck_10 = None
for topic, msg, t in bag.read_messages("/tf"):
if msg.transforms[0].child_frame_id == "Puck":
if i == 0:
tf_puck_start = msg.transforms[0]
if i == 10:
tf_puck_10 = msg.transforms[0]
break
i += 1
if msg.transforms[0].child_frame_id == "Table":
if j == 0:
tf_table = msg.transforms[0]
j += 1
puck_start_T = tf2_geometry_msgs.transform_to_kdl(tf_table).Inverse() * \
tf2_geometry_msgs.transform_to_kdl(tf_puck_start)
puck_10_T = tf2_geometry_msgs.transform_to_kdl(tf_table).Inverse() * \
tf2_geometry_msgs.transform_to_kdl(tf_puck_10)
_, _, yaw_start = puck_start_T.M.GetRPY()
_, _, yaw_10 = puck_10_T.M.GetRPY()
t_start = tf_puck_start.header.stamp.to_sec()
t_10 = tf_puck_10.header.stamp.to_sec()
p_start = np.array([puck_start_T.p.x(), puck_start_T.p.y()])
p_10 = np.array([puck_10_T.p.x(), puck_10_T.p.y()])
lin_vel_start = (p_10 - p_start) / (t_10 - t_start)
ang_vel_start = angles.shortest_angular_distance(yaw_start,
yaw_10) / (t_10 - t_start)
return np.hstack([p_start,
yaw_start]), np.hstack([lin_vel_start,
ang_vel_start]), tf_table
def play_images(bag_file: Bag, topics: list) -> None:
"""
Play the data in a bag file.
Args:
bag_file: the bag file to play
topics: the list of topics to play
Returns:
None
"""
# open windows to stream the camera and a priori image data to
windows = {topic: None for topic in topics}
# iterate over the messages
progress = tqdm(total=bag_file.get_message_count(topic_filters=topics))
for topic, msg, time in bag_file.read_messages(topics=topics):
# if topic is camera, unwrap and send to the camera window
if topic in topics:
# update the progress bar with an iteration
progress.update(1)
# update the progress with a post fix
progress.set_postfix(time=time)
# if the camera window isn't open, open it
if windows[topic] is None:
title = '{} ({})'.format(bag_file.filename, topic)
windows[topic] = Window(title, msg.height, msg.width)
# get the pixels of the camera image and display them
img = get_camera_image(msg.data, windows[topic].shape)
if msg.encoding == 'bgr8':
img = img[..., ::-1]
windows[topic].show(img[..., :3])
# shut down the viewer windows
for window in windows.values():
if window is not None:
window.close()
outimgs_path = '/home/kaue/data/extracted_images/newextracted'
from geometry_msgs.msg import _TwistStamped
t1 = time.time()
# bagpath = "/home/kauevestena/data/rosbags/2019-07-11-16-21-46.bag"
bagpath = msc.joinToHome('data/rosbags/2019-07-11-16-21-46.bag')
print('opened',bagpath)
current_bag = Bag(bagpath)
# print(current_bag.get_type_and_topic_info())
messages = current_bag.read_messages()
t0 = messages.next().timestamp.to_sec()
topic_names = {
"imu" : "imu/data",
"img" : "/raspicam_node/image/compressed",
"gnss" : "/fix"
}
# imu_list = []
# img_list = []
# gnss_list = []
# # # # the big dict that will store all the relevant data
pickle_outpath = os.path.join(msc.PICKLESPATH,msc.filenameFromPathWtithoutExt(bagpath)+'_idx.pickle')
def on_enter(self, userdata):
self._failed = False
try:
# Initialization
# userdata
self._bag_filename = userdata.bag_filename
self._trajectories_command = userdata.trajectories_command
# joint definitions
l_arm_range = range(16, 23)
r_arm_range = range(23, 30)
atlasJointNames = [
'back_bkz', 'back_bky', 'back_bkx', 'neck_ry', 'l_leg_hpz',
'l_leg_hpx', 'l_leg_hpy', 'l_leg_kny', 'l_leg_aky',
'l_leg_akx', 'r_leg_hpz', 'r_leg_hpx', 'r_leg_hpy',
'r_leg_kny', 'r_leg_aky', 'r_leg_akx', 'l_arm_shz',
'l_arm_shx', 'l_arm_ely', 'l_arm_elx', 'l_arm_wry',
'l_arm_wrx', 'l_arm_wry2', 'r_arm_shz', 'r_arm_shx',
'r_arm_ely', 'r_arm_elx', 'r_arm_wry', 'r_arm_wrx',
'r_arm_wry2'
]
joint_position_cmd = [0] * len(l_arm_range)
# starting and stopping time of measurement data to be taken into account for cailbration
time_start = 0
time_end = 0
time_tf = 0
# gravity vector in world frame
g_world = np.resize(np.array([0, 0, -9.81]), (3, 1))
# take this number of data points for each pose.
max_data_points_per_pose = 100
# rotation of the force torque sensor in world frame (quaternions)
ft_rotation = [0, 0, 0, 0]
# loop over all ft sensors to calibrate
for chain in self._calibration_chain:
tf_chain = ['/pelvis', 'ltorso', 'mtorso', 'utorso']
if chain == 'left_arm':
joint_range = l_arm_range
tfname = 'l_hand'
tf_chain.extend([
'l_clav', 'l_scap', 'l_uarm', 'l_larm', 'l_ufarm',
'l_lfarm', 'l_hand'
])
elif chain == 'right_arm':
joint_range = r_arm_range
tf_chain.extend([
'r_clav', 'r_scap', 'r_uarm', 'r_larm', 'r_ufarm',
'r_lfarm', 'r_hand'
])
tfname = 'r_hand'
else:
Logger.logwarn(
'CalculateForceTorqueCalibration: Undefined chain %s',
chain)
# initialize transformations
tf_data = [
] # frame to frame transformation with numerical indexes from the tf_chain
tf_data_cum = [
] # world to frame transformation with numerical indexes from the tf_chain
for i in range(len(tf_chain)):
tf_data.append([0, 0, 0, 0])
tf_data_cum.append([0, 0, 0, 0])
# get number of poses from the commanded trajectories. 2 Points per Pose
number_of_poses = len(
self._trajectories_command[chain].points) / 2
# define information matrix for calibration
InfMat = np.zeros(
(6 * max_data_points_per_pose * number_of_poses, 10))
InfMat_i = 0 # Index in this Matrix
# definie measurement vector for calibration
MeasVec = np.zeros(
(6 * max_data_points_per_pose * number_of_poses, 1))
# commanded joint trajectory for current chain
# all commanded positions in this trajectory
current_traj_cmd = self._trajectories_command[chain]
# flag if the time interval was already set:
timesetflag = False
point_index_cmd = 0 # start with index 0
transformation_available = False
# read time series from bag file
# check in desired input trajectory, at which time a position is reached
# take the data in a defined time period after the commanded new position
bag_from_robot = Bag(os.path.expanduser(
self._bag_filename)) # open bag file
Logger.loginfo(
'CalculateForceTorqueCalibration: Calibrate %s from %s. Using %d different positions from trajectory command. Expecting 2 commanded positions per pose'
%
(chain, self._bag_filename, len(current_traj_cmd.points)))
for topic, msg, t in bag_from_robot.read_messages(topics=[
'/flor/controller/atlas_state', '/tf',
'/flor/controller/joint_command'
]):
# Check if all desired poses have been reached
if point_index_cmd > len(current_traj_cmd.points) - 1:
break
current_position_cmd = current_traj_cmd.points[
point_index_cmd].positions
####################################################
# Check tf message: remember last transformation from world to ft sensor
if topic == '/tf':
# loop over all transformations inside this tf message
for j in range(len(msg.transforms)):
data = msg.transforms[j]
tr = data.transform # current transformation
header = data.header
# check if frame matches on of the saved frames
for i in range(len(tf_chain)):
if data.child_frame_id == tf_chain[i]:
tf_data[i] = np.array([
tr.rotation.x, tr.rotation.y,
tr.rotation.z, tr.rotation.w
])
# calculate the newest transformation to the force torque sensor
tf_data_cum[0] = tf_data[0]
for i in range(1, len(tf_chain)):
tf_data_cum[
i] = q = tf.transformations.quaternion_multiply(
tf_data_cum[i - 1], tf_data[i])
if i == len(tf_chain) - 1:
time_tf = msg.transforms[
0].header.stamp.to_sec()
# nothing to do with tf message
# check if all transformations are available. (up to the last link)
if np.any(
tf_data_cum[-1]
): # real part of quaternion unequal to zero: data exists
transformation_available = True
continue
else:
time_msg = msg.header.stamp.to_sec()
####################################################
# check if timestamp is interesting. Then check the
if not (time_msg > time_start
and time_msg < time_end) and not timesetflag:
# the timestamp is not in evaluation interval. Look for reaching of the pose
# record data (see below)
if topic == '/flor/controller/joint_command':
# get the commanded position
for i in range(len(joint_range)):
osrf_ndx = joint_range[i]
joint_position_cmd[i] = msg.position[osrf_ndx]
# Check if position command matches the currently expected commanded position
pos_reached = True
for i in range(len(joint_range)):
if abs(joint_position_cmd[i] -
current_position_cmd[i]) > 0.001:
pos_reached = False
break
if pos_reached:
# end time of the values regarded for calibration data: take the latest time possible for this pose
time_end = msg.header.stamp.to_sec(
) + self._settlingtime
# starting time for calibration: Take 100ms
time_start = time_end - 0.1
data_points_for_this_pose = 0
timesetflag = True
# take the next point next time. Each pose consists of two trajectory points (one for reaching, one for settling).
# take the second one
point_index_cmd = point_index_cmd + 2
if not pos_reached:
# The commanded position has not been reached. Skip
continue
continue # continue here, because data aquisition is triggered by the time_start, time_end
if time_msg > time_end:
timesetflag = False # prepare for new evaluation interval
####################################################
# check if enough datapoints for this pose have been collected
if data_points_for_this_pose > max_data_points_per_pose:
# already enough data points for this pose
continue
####################################################
# Check if message is atlas_state
# IF this is the case, fill information matrix
if topic != '/flor/controller/atlas_state':
continue
# Extract measured force and torque
if chain == 'left_arm':
FT = [
msg.l_hand.force.x, msg.l_hand.force.y,
msg.l_hand.force.z, msg.l_hand.torque.x,
msg.l_hand.torque.y, msg.l_hand.torque.z
]
elif chain == 'right_arm':
FT = [
msg.r_hand.force.x, msg.r_hand.force.y,
msg.r_hand.force.z, msg.r_hand.torque.x,
msg.r_hand.torque.y, msg.r_hand.torque.z
]
# calculate gravitation vector
if not transformation_available:
Logger.logwarn(
'No tf messages available at time %1.4f.' %
time_msg)
continue
R = tf.transformations.quaternion_matrix(tf_data_cum[-1])
g = np.dot((R[0:3, 0:3]).transpose(), g_world)
gx = g[0]
gy = g[1]
gz = g[2]
# fill information matrix for this data point. Source: [1], equ. (7)
M = np.zeros((6, 10))
M[0, 0] = gx
M[1, 0] = gy
M[2, 0] = gz
M[3, 2] = gz
M[3, 3] = -gy
M[4, 1] = -gz
M[4, 3] = gx
M[5, 1] = gy
M[5, 2] = -gx
M[0, 4] = 1.0
M[1, 5] = 1.0
M[2, 6] = 1.0
M[3, 7] = 1.0
M[4, 8] = 1.0
M[5, 9] = 1.0
# fill big information matrix and vector (stack small information matrizes)
for i in range(6):
for j in range(10):
InfMat[InfMat_i * 6 + i, j] = M[i, j]
MeasVec[InfMat_i * 6 + i, 0] = FT[i]
InfMat_i = InfMat_i + 1 # increase index
data_points_for_this_pose = data_points_for_this_pose + 1
# shorten big information matrix
if InfMat_i < max_data_points_per_pose * number_of_poses:
InfMat_calc = InfMat[0:(6 * InfMat_i) - 1, :]
MeasVec_calc = MeasVec[0:(6 * InfMat_i - 1), :]
else:
InfMat_calc = InfMat
MeasVec_calc = MeasVec
# calculate calibration data
if chain in self._static_calibration_data.keys(
): # calculate calibration with given static parameters
# bring colums with first parameters on the other side of the equation
if len(self._static_calibration_data[chain]) != 4:
Logger.logwarn(
"CalculateForceTorqueCalibration: Given static calibration data for %s has length %d. Required 4 entries."
% (chain, len(k_fix)))
self._failed = True
return
# convert physical parameters to identification parameters (mass, 1st moment)
m = self._static_calibration_data[chain][0]
if m == 0:
mom_x = 0
mom_y = 0
mom_z = 0
elif m > 0:
mom_x = self._static_calibration_data[chain][1] / m
mom_y = self._static_calibration_data[chain][2] / m
mom_z = self._static_calibration_data[chain][3] / m
else:
Logger.logwarn(
"CalculateForceTorqueCalibration: Negative mass (%f) for calibration requested. Abort."
% mass)
self._failed = True
return
k_fix = np.resize(np.array([m, mom_x, mom_y, mom_z]),
(4, 1))
Logger.loginfo(
"CalculateForceTorqueCalibration:static calibration data for %s given: %s. Reduce equation system"
% (chain, str(self._static_calibration_data[chain])))
MeasVec_calc_corr = np.subtract(
np.array(MeasVec_calc),
np.dot(InfMat_calc[:, 0:4], k_fix))
InfMat_calc_corr = InfMat_calc[:, 4:
10] # only the last 6 colums which correspond to the sensor offsets
k = np.linalg.lstsq(
InfMat_calc_corr,
MeasVec_calc_corr)[0] # solve reduced equation system
k_calibration = self._static_calibration_data[chain]
k_calibration.extend(k)
else: # calculate normally with all parameters unknown
k = np.linalg.lstsq(InfMat_calc, MeasVec_calc)[0]
# convert to physical parameters (first moment -> mass)
if k[0] > 0:
k_calibration = [
k[0], k[1] / k[0], k[2] / k[0], k[3] / k[0], k[4],
k[5], k[6], k[7], k[8], k[9]
]
else:
Logger.loginfo(
"CalculateForceTorqueCalibration:Calibration brought negative mass %f"
% k[0])
k_calibration = [
0.0, 0.0, 0.0, 0.0, k[4], k[5], k[6], k[7], k[8],
k[9]
]
Logger.loginfo(
"CalculateForceTorqueCalibration:calibration data for %s" %
chain)
Logger.loginfo("CalculateForceTorqueCalibration:mass: %f" %
float(k_calibration[0]))
Logger.loginfo(
"CalculateForceTorqueCalibration:center of mass: %f %f %f"
% (k_calibration[1], k_calibration[2], k_calibration[3]))
Logger.loginfo(
"CalculateForceTorqueCalibration:F offset: %f %f %f" %
(k_calibration[4], k_calibration[5], k_calibration[6]))
Logger.loginfo(
"CalculateForceTorqueCalibration:M offset: %f %f %f" %
(k_calibration[7], k_calibration[8], k_calibration[9]))
self._ft_calib_data[chain] = k_calibration
bag_from_robot.close()
userdata.ft_calib_data = self._ft_calib_data
Logger.loginfo(
'CalculateForceTorqueCalibration:Calibration finished')
self._done = True
except Exception as e:
Logger.logwarn(
'CalculateForceTorqueCalibration:Unable to calculate calibration:\n%s'
% str(e))
self._failed = True
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 on_enter(self, userdata):
self._failed = False
try:
# Initialization
# userdata
self._bag_filename = userdata.bag_filename
self._trajectories_command = userdata.trajectories_command
# joint definitions
l_arm_range = range(16,23);
r_arm_range = range(23,30);
atlasJointNames = [
'back_bkz', 'back_bky', 'back_bkx', 'neck_ry',
'l_leg_hpz', 'l_leg_hpx', 'l_leg_hpy', 'l_leg_kny', 'l_leg_aky', 'l_leg_akx',
'r_leg_hpz', 'r_leg_hpx', 'r_leg_hpy', 'r_leg_kny', 'r_leg_aky', 'r_leg_akx',
'l_arm_shz', 'l_arm_shx', 'l_arm_ely', 'l_arm_elx', 'l_arm_wry', 'l_arm_wrx', 'l_arm_wry2',
'r_arm_shz', 'r_arm_shx', 'r_arm_ely', 'r_arm_elx', 'r_arm_wry', 'r_arm_wrx', 'r_arm_wry2']
joint_position_cmd = [0]*len(l_arm_range)
# starting and stopping time of measurement data to be taken into account for cailbration
time_start = 0
time_end = 0
time_tf = 0
# gravity vector in world frame
g_world = np.resize(np.array([0,0,-9.81]),(3,1))
# take this number of data points for each pose.
max_data_points_per_pose = 100
# rotation of the force torque sensor in world frame (quaternions)
ft_rotation = [0,0,0,0]
# loop over all ft sensors to calibrate
for chain in self._calibration_chain:
tf_chain=['/pelvis', 'ltorso', 'mtorso', 'utorso']
if chain == 'left_arm':
joint_range = l_arm_range
tfname = 'l_hand'
tf_chain.extend(['l_clav', 'l_scap', 'l_uarm', 'l_larm', 'l_ufarm', 'l_lfarm', 'l_hand'])
elif chain == 'right_arm':
joint_range = r_arm_range
tf_chain.extend(['r_clav', 'r_scap', 'r_uarm', 'r_larm', 'r_ufarm', 'r_lfarm', 'r_hand'])
tfname = 'r_hand'
else:
Logger.logwarn('CalculateForceTorqueCalibration: Undefined chain %s', chain)
# initialize transformations
tf_data = [] # frame to frame transformation with numerical indexes from the tf_chain
tf_data_cum = [] # world to frame transformation with numerical indexes from the tf_chain
for i in range(len(tf_chain)):
tf_data.append([0,0,0,0])
tf_data_cum.append([0,0,0,0])
# get number of poses from the commanded trajectories. 2 Points per Pose
number_of_poses = len(self._trajectories_command[chain].points) / 2
# define information matrix for calibration
InfMat = np.zeros((6*max_data_points_per_pose*number_of_poses, 10))
InfMat_i = 0 # Index in this Matrix
# definie measurement vector for calibration
MeasVec = np.zeros((6*max_data_points_per_pose*number_of_poses, 1))
# commanded joint trajectory for current chain
# all commanded positions in this trajectory
current_traj_cmd = self._trajectories_command[chain]
# flag if the time interval was already set:
timesetflag = False
point_index_cmd = 0 # start with index 0
transformation_available = False
# read time series from bag file
# check in desired input trajectory, at which time a position is reached
# take the data in a defined time period after the commanded new position
bag_from_robot = Bag(os.path.expanduser(self._bag_filename)) # open bag file
Logger.loginfo('CalculateForceTorqueCalibration: Calibrate %s from %s. Using %d different positions from trajectory command. Expecting 2 commanded positions per pose' % (chain, self._bag_filename, len(current_traj_cmd.points)) )
for topic, msg, t in bag_from_robot.read_messages(topics=['/flor/controller/atlas_state', '/tf', '/flor/controller/joint_command']):
# Check if all desired poses have been reached
if point_index_cmd > len(current_traj_cmd.points)-1:
break
current_position_cmd = current_traj_cmd.points[point_index_cmd].positions
####################################################
# Check tf message: remember last transformation from world to ft sensor
if topic == '/tf':
# loop over all transformations inside this tf message
for j in range(len(msg.transforms)):
data = msg.transforms[j]
tr = data.transform # current transformation
header = data.header
# check if frame matches on of the saved frames
for i in range(len(tf_chain)):
if data.child_frame_id == tf_chain[i]:
tf_data[i] = np.array([tr.rotation.x, tr.rotation.y, tr.rotation.z, tr.rotation.w])
# calculate the newest transformation to the force torque sensor
tf_data_cum[0] = tf_data[0]
for i in range(1,len(tf_chain)):
tf_data_cum[i] = q = tf.transformations.quaternion_multiply(tf_data_cum[i-1], tf_data[i])
if i == len(tf_chain)-1:
time_tf = msg.transforms[0].header.stamp.to_sec()
# nothing to do with tf message
# check if all transformations are available. (up to the last link)
if np.any(tf_data_cum[-1]): # real part of quaternion unequal to zero: data exists
transformation_available = True
continue
else:
time_msg = msg.header.stamp.to_sec()
####################################################
# check if timestamp is interesting. Then check the
if not (time_msg > time_start and time_msg < time_end) and not timesetflag:
# the timestamp is not in evaluation interval. Look for reaching of the pose
# record data (see below)
if topic == '/flor/controller/joint_command':
# get the commanded position
for i in range(len(joint_range)):
osrf_ndx = joint_range[i]
joint_position_cmd[i] = msg.position[osrf_ndx]
# Check if position command matches the currently expected commanded position
pos_reached = True
for i in range(len(joint_range)):
if abs(joint_position_cmd[i]-current_position_cmd[i]) > 0.001:
pos_reached = False
break
if pos_reached:
# end time of the values regarded for calibration data: take the latest time possible for this pose
time_end = msg.header.stamp.to_sec() + self._settlingtime
# starting time for calibration: Take 100ms
time_start = time_end - 0.1
data_points_for_this_pose = 0
timesetflag = True
# take the next point next time. Each pose consists of two trajectory points (one for reaching, one for settling).
# take the second one
point_index_cmd = point_index_cmd + 2
if not pos_reached:
# The commanded position has not been reached. Skip
continue
continue # continue here, because data aquisition is triggered by the time_start, time_end
if time_msg > time_end:
timesetflag = False # prepare for new evaluation interval
####################################################
# check if enough datapoints for this pose have been collected
if data_points_for_this_pose > max_data_points_per_pose:
# already enough data points for this pose
continue
####################################################
# Check if message is atlas_state
# IF this is the case, fill information matrix
if topic != '/flor/controller/atlas_state':
continue
# Extract measured force and torque
if chain == 'left_arm':
FT = [msg.l_hand.force.x, msg.l_hand.force.y, msg.l_hand.force.z, msg.l_hand.torque.x, msg.l_hand.torque.y, msg.l_hand.torque.z ]
elif chain == 'right_arm':
FT = [msg.r_hand.force.x, msg.r_hand.force.y, msg.r_hand.force.z, msg.r_hand.torque.x, msg.r_hand.torque.y, msg.r_hand.torque.z ]
# calculate gravitation vector
if not transformation_available:
Logger.logwarn('No tf messages available at time %1.4f.' % time_msg)
continue
R = tf.transformations.quaternion_matrix(tf_data_cum[-1])
g = np.dot((R[0:3,0:3]).transpose(), g_world)
gx = g[0]
gy = g[1]
gz = g[2]
# fill information matrix for this data point. Source: [1], equ. (7)
M = np.zeros((6, 10))
M[0,0] = gx
M[1,0] = gy
M[2,0] = gz
M[3,2] = gz
M[3,3] = -gy
M[4,1] = -gz
M[4,3] = gx
M[5,1] = gy
M[5,2] = -gx
M[0,4] = 1.0
M[1,5] = 1.0
M[2,6] = 1.0
M[3,7] = 1.0
M[4,8] = 1.0
M[5,9] = 1.0
# fill big information matrix and vector (stack small information matrizes)
for i in range(6):
for j in range(10):
InfMat[InfMat_i*6+i,j] = M[i,j]
MeasVec[InfMat_i*6+i,0] = FT[i]
InfMat_i = InfMat_i + 1 # increase index
data_points_for_this_pose = data_points_for_this_pose + 1
# shorten big information matrix
if InfMat_i < max_data_points_per_pose*number_of_poses:
InfMat_calc = InfMat[0:(6*InfMat_i)-1,:]
MeasVec_calc = MeasVec[0:(6*InfMat_i-1),:]
else:
InfMat_calc = InfMat
MeasVec_calc = MeasVec
# calculate calibration data
if chain in self._static_calibration_data.keys(): # calculate calibration with given static parameters
# bring colums with first parameters on the other side of the equation
if len(self._static_calibration_data[chain]) != 4:
Logger.logwarn( "CalculateForceTorqueCalibration: Given static calibration data for %s has length %d. Required 4 entries." % (chain, len(k_fix)) )
self._failed = True
return
# convert physical parameters to identification parameters (mass, 1st moment)
m = self._static_calibration_data[chain][0]
if m == 0:
mom_x = 0
mom_y = 0
mom_z = 0
elif m > 0:
mom_x = self._static_calibration_data[chain][1]/m
mom_y = self._static_calibration_data[chain][2]/m
mom_z = self._static_calibration_data[chain][3]/m
else:
Logger.logwarn( "CalculateForceTorqueCalibration: Negative mass (%f) for calibration requested. Abort." % mass )
self._failed = True
return
k_fix = np.resize(np.array([m, mom_x, mom_y, mom_z]),(4,1))
Logger.loginfo( "CalculateForceTorqueCalibration:static calibration data for %s given: %s. Reduce equation system" % (chain, str(self._static_calibration_data[chain])) )
MeasVec_calc_corr = np.subtract(np.array(MeasVec_calc), np.dot(InfMat_calc[:,0:4], k_fix))
InfMat_calc_corr = InfMat_calc[:,4:10] # only the last 6 colums which correspond to the sensor offsets
k = np.linalg.lstsq(InfMat_calc_corr, MeasVec_calc_corr)[0] # solve reduced equation system
k_calibration = self._static_calibration_data[chain]
k_calibration.extend(k)
else: # calculate normally with all parameters unknown
k = np.linalg.lstsq(InfMat_calc, MeasVec_calc)[0]
# convert to physical parameters (first moment -> mass)
if k[0] > 0:
k_calibration = [k[0], k[1]/k[0], k[2]/k[0], k[3]/k[0], k[4], k[5], k[6], k[7], k[8], k[9]]
else:
Logger.loginfo("CalculateForceTorqueCalibration:Calibration brought negative mass %f" % k[0])
k_calibration = [0.0, 0.0, 0.0, 0.0, k[4], k[5], k[6], k[7], k[8], k[9]]
Logger.loginfo("CalculateForceTorqueCalibration:calibration data for %s" % chain)
Logger.loginfo("CalculateForceTorqueCalibration:mass: %f" % float(k_calibration[0]))
Logger.loginfo("CalculateForceTorqueCalibration:center of mass: %f %f %f" % (k_calibration[1], k_calibration[2], k_calibration[3]))
Logger.loginfo("CalculateForceTorqueCalibration:F offset: %f %f %f" % (k_calibration[4], k_calibration[5], k_calibration[6]))
Logger.loginfo("CalculateForceTorqueCalibration:M offset: %f %f %f" % (k_calibration[7], k_calibration[8], k_calibration[9]))
self._ft_calib_data[chain] = k_calibration
bag_from_robot.close()
userdata.ft_calib_data = self._ft_calib_data
Logger.loginfo('CalculateForceTorqueCalibration:Calibration finished')
self._done = True
except Exception as e:
Logger.logwarn('CalculateForceTorqueCalibration:Unable to calculate calibration:\n%s' % str(e))
self._failed = True
class Rosbag(ForeignDataWrapper):
def __init__(self, options, columns=None):
super(Rosbag, self).__init__(options, columns)
Bag = import_bag(options)
self.filename = options.pop('rosbag_path', "") + options.pop('rosbag')
self.topic = options.pop('topic', None)
pointcloud_formats = strtobool(options.pop('metadata', 'false'))
self.patch_column = options.pop('patch_column', 'points').strip()
self.patch_columns = options.pop('patch_columns', '*').strip()
self.patch_columns = [
col.strip() for col in self.patch_columns.split(',')
if col.strip()
]
self.patch_count_default = int(options.pop('patch_count_default',
1000))
# 0 => 1 patch per message
self.patch_count_pointcloud = int(
options.pop('patch_count_pointcloud', 0))
assert (self.patch_count_default > 0)
assert (self.patch_count_pointcloud >= 0)
self.pcid = int(options.pop('pcid', 0))
self.bag = Bag(self.filename, 'r')
self.topics = self.bag.get_type_and_topic_info().topics
self.pointcloud_formats = None
if pointcloud_formats:
self.pointcloud_formats = []
topics = self.topic.split(',') if self.topic else self.topics
for i, topic in enumerate(self.topics):
if topic not in topics:
continue
infos = self.topics[topic]
columns, patch_schema, patch_ply_header, _, patch_columns, patch_srid = \
get_columns(self.bag, topic, infos, self.pcid+i+1, self.patch_column,
self.patch_columns)
self.pointcloud_formats.append({
'pcid':
self.pcid + i + 1,
'srid':
patch_srid,
'schema':
patch_schema,
'format':
columns[self.patch_column][3],
'rostype':
infos.msg_type,
'columns':
patch_columns,
'ply_header':
patch_ply_header,
})
return
if not self.topic:
log_to_postgres('"topic" option is required', ERROR)
self.infos = self.topics[self.topic]
(self.columns, self.patch_schema, self.patch_ply_header, self.endianness,
self.patch_columns, self.patch_srid) = \
get_columns(self.bag, self.topic, self.infos, self.pcid, self.patch_column,
self.patch_columns)
if columns:
missing = set(columns) - set(self.columns.keys())
columns = list(c for c in self.columns.keys() if c in columns)
if missing:
missing = ", ".join(sorted(missing))
support = ", ".join(sorted(self.columns.keys()))
log_to_postgres(
"extra unsupported columns : {}".format(missing),
WARNING,
hint="supported columns : {}".format(support))
self.columns = {col: self.columns[col] for col in columns}
if options:
log_to_postgres(
"extra unsupported options : {}".format(options.keys()),
WARNING)
@classmethod
def import_schema(self, schema, srv_options, options, restriction_type,
restricts):
Bag = import_bag(srv_options)
pcid_str = options.pop('pcid', srv_options.pop('pcid', 0))
pcid = int(pcid_str)
patch_column = options.pop('patch_column',
srv_options.pop('patch_column', 'points'))
patch_columns = options.pop('patch_columns', '*').strip()
patch_columns = [
col.strip() for col in patch_columns.split(',') if col.strip()
]
filename = srv_options.pop('rosbag_path', "") + options.pop(
'rosbag_path', "") + schema
bag = Bag(filename, 'r')
tablecols = []
topics = bag.get_type_and_topic_info().topics
pcid_for_topic = {k: pcid + 1 + i for i, k in enumerate(topics.keys())}
pointcloud_formats = True
if restriction_type is 'limit':
topics = {k: v for k, v in topics.items() if k in restricts}
pointcloud_formats = 'pointcloud_formats' in restricts
elif restriction_type is 'except':
topics = {k: v for k, v in topics.items() if k not in restricts}
pointcloud_formats = 'pointcloud_formats' not in restricts
tabledefs = []
if pointcloud_formats:
tablecols = [
ColumnDefinition('pcid', type_name='integer'),
ColumnDefinition('srid', type_name='integer'),
ColumnDefinition('schema', type_name='text'),
ColumnDefinition('format', type_name='text'),
ColumnDefinition('rostype', type_name='text'),
ColumnDefinition('columns', type_name='text[]'),
ColumnDefinition('ply_header', type_name='text'),
]
tableopts = {
'metadata': 'true',
'rosbag': schema,
'pcid': pcid_str
}
tabledefs.append(
TableDefinition("pointcloud_formats",
columns=tablecols,
options=tableopts))
for topic, infos in topics.items():
pcid = pcid_for_topic[topic]
columns, _, _, _, _, _ = get_columns(bag, topic, infos, pcid,
patch_column, patch_columns)
tablecols = [get_column_def(k, *v) for k, v in columns.items()]
tableopts = {'topic': topic, 'rosbag': schema, 'pcid': str(pcid)}
tabledefs.append(
TableDefinition(topic, columns=tablecols, options=tableopts))
return tabledefs
def execute(self, quals, columns):
if self.pointcloud_formats is not None:
for f in self.pointcloud_formats:
yield f
return
self.patch_data = ''
from rospy.rostime import Time
tmin = None
tmax = None
for qual in quals:
if qual.field_name == "time":
t = int(qual.value)
t = Time(t / 1000000000, t % 1000000000)
if qual.operator in ['=', '>', '>=']:
tmin = t
if qual.operator in ['=', '<', '<=']:
tmax = t
for topic, msg, t in self.bag.read_messages(topics=self.topic,
start_time=tmin,
end_time=tmax):
for row in self.get_rows(topic, msg, t, columns):
yield row
# flush leftover patch data
if self.patch_data and self.last_row:
count = int((len(self.patch_data) / self.point_size))
# in replicating mode, a single leftover point must not be reported
if count > 1 or self.patch_step_size == self.patch_size:
res = self.last_row
if self.patch_column in columns:
res[self.patch_column] = hexlify(
pack('=b3I', self.endianness, self.pcid, 0, count) +
self.patch_data)
if self.patch_ply_header and 'ply' in columns:
self.ply_info['count'] = count
res['ply'] = self.patch_ply_header.format(
**self.ply_info) + self.patch_data
yield res
def get_rows(self, topic, msg, t, columns, toplevel=True):
if toplevel and len(msg.__slots__) == 1:
attr = getattr(msg, msg.__slots__[0])
if isinstance(attr, list):
for msg in attr:
for row in self.get_rows(topic, msg, t, columns, False):
yield row
return
res = {}
data_columns = set(columns)
if self.patch_column in columns:
data_columns = data_columns.union(self.patch_columns) - set(
[self.patch_column])
if "filename" in data_columns:
res["filename"] = self.filename
if "topic" in data_columns:
res["topic"] = topic
if "time" in data_columns:
res["time"] = t.to_nsec()
if self.infos.msg_type == 'sensor_msgs/PointCloud2':
self.patch_count = self.patch_count_pointcloud or (msg.width *
msg.height)
self.point_size = msg.point_step
self.patch_size = self.patch_count * self.point_size
self.patch_step_size = self.patch_size
self.endianness = 0 if msg.is_bigendian else 1
data_columns = data_columns - set(['ply', self.patch_column])
self.patch_data += msg.data
data_columns = data_columns - set(res.keys())
for column in data_columns:
attr = msg
for col in column.split('.'):
if isinstance(attr, list):
attr = tuple(getattr(a, col) for a in attr)
else:
attr = getattr(attr, col)
if hasattr(attr, "to_nsec"):
attr = attr.to_nsec()
elif hasattr(attr, "x"):
if hasattr(attr, "w"):
attr = (attr.x, attr.y, attr.z, attr.w)
else:
attr = (attr.x, attr.y, attr.z)
elif isinstance(attr, str):
fmt = self.columns[column][3]
if fmt:
attr = unpack(fmt, attr)
res[column] = attr
if self.patch_column in columns and not self.infos.msg_type == 'sensor_msgs/PointCloud2':
fmt = self.columns[self.patch_column][3]
self.patch_count = self.patch_count_default
self.point_size = calcsize(fmt)
self.patch_size = self.patch_count * self.point_size
self.patch_step_size = self.patch_size - self.point_size
self.patch_data += get_point_data(res, self.patch_columns, fmt)
res = {k: v for k, v in res.items() if k not in self.patch_columns}
if not self.patch_data:
yield res
else:
# todo: ensure current res and previous res are equal if there is some leftover
# patch_data
while len(self.patch_data) >= self.patch_size:
data = self.patch_data[0:self.patch_size]
count = int(self.patch_size / self.point_size)
res[self.patch_column] = hexlify(
pack('=b3I', self.endianness, self.pcid, 0, count) + data)
if self.patch_ply_header and 'ply' in columns:
self.ply_info = {
'endianness': 'big' if self.endianness else 'little',
'filename': self.filename,
'topic': self.topic,
'count': count
}
res['ply'] = self.patch_ply_header.format(
**self.ply_info) + data
self.patch_data = self.patch_data[self.patch_step_size:]
yield res
self.last_row = res
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)
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()
class Rosbag(ForeignDataWrapper):
def __init__(self, options, columns=None):
super(Rosbag, self).__init__(options, columns)
Bag = import_bag(options)
self.filename = options.pop('rosbag_path', "") + options.pop('rosbag')
self.topic = options.pop('topic', None)
pointcloud_formats = strtobool(options.pop('metadata', 'false'))
self.patch_column = options.pop('patch_column', 'points').strip()
self.patch_columns = options.pop('patch_columns', '*').strip()
self.patch_columns = [col.strip() for col in self.patch_columns.split(',') if col.strip()]
self.patch_count_default = int(options.pop('patch_count_default', 1000))
# 0 => 1 patch per message
self.patch_count_pointcloud = int(options.pop('patch_count_pointcloud', 0))
assert(self.patch_count_default > 0)
assert(self.patch_count_pointcloud >= 0)
self.pcid = int(options.pop('pcid', 0))
self.bag = Bag(self.filename, 'r')
self.topics = self.bag.get_type_and_topic_info().topics
self.pointcloud_formats = None
if pointcloud_formats:
self.pointcloud_formats = []
topics = self.topic.split(',') if self.topic else self.topics
for i, topic in enumerate(self.topics):
if topic not in topics:
continue
infos = self.topics[topic]
columns, patch_schema, patch_ply_header, _, patch_columns, patch_srid = \
get_columns(self.bag, topic, infos, self.pcid+i+1, self.patch_column,
self.patch_columns)
self.pointcloud_formats.append({
'pcid': self.pcid+i+1,
'srid': patch_srid,
'schema': patch_schema,
'format': columns[self.patch_column][3],
'rostype': infos.msg_type,
'columns': patch_columns,
'ply_header': patch_ply_header,
})
return
if not self.topic:
log_to_postgres('"topic" option is required', ERROR)
self.infos = self.topics[self.topic]
(self.columns, self.patch_schema, self.patch_ply_header, self.endianness,
self.patch_columns, self.patch_srid) = \
get_columns(self.bag, self.topic, self.infos, self.pcid, self.patch_column,
self.patch_columns)
if columns:
missing = set(columns) - set(self.columns.keys())
columns = list(c for c in self.columns.keys() if c in columns)
if missing:
missing = ", ".join(sorted(missing))
support = ", ".join(sorted(self.columns.keys()))
log_to_postgres(
"extra unsupported columns : {}".format(missing), WARNING,
hint="supported columns : {}".format(support))
self.columns = {col: self.columns[col] for col in columns}
if options:
log_to_postgres("extra unsupported options : {}".format(
options.keys()), WARNING)
@classmethod
def import_schema(self, schema, srv_options, options,
restriction_type, restricts):
Bag = import_bag(srv_options)
pcid_str = options.pop('pcid', srv_options.pop('pcid', 0))
pcid = int(pcid_str)
patch_column = options.pop('patch_column', srv_options.pop('patch_column', 'points'))
patch_columns = options.pop('patch_columns', '*').strip()
patch_columns = [col.strip() for col in patch_columns.split(',') if col.strip()]
filename = srv_options.pop('rosbag_path', "") + options.pop('rosbag_path', "") + schema
bag = Bag(filename, 'r')
tablecols = []
topics = bag.get_type_and_topic_info().topics
pcid_for_topic = {k: pcid+1+i for i, k in enumerate(topics.keys())}
pointcloud_formats = True
if restriction_type is 'limit':
topics = {k: v for k, v in topics.items() if k in restricts}
pointcloud_formats = 'pointcloud_formats' in restricts
elif restriction_type is 'except':
topics = {k: v for k, v in topics.items() if k not in restricts}
pointcloud_formats = 'pointcloud_formats' not in restricts
tabledefs = []
if pointcloud_formats:
tablecols = [
ColumnDefinition('pcid', type_name='integer'),
ColumnDefinition('srid', type_name='integer'),
ColumnDefinition('schema', type_name='text'),
ColumnDefinition('format', type_name='text'),
ColumnDefinition('rostype', type_name='text'),
ColumnDefinition('columns', type_name='text[]'),
ColumnDefinition('ply_header', type_name='text'),
]
tableopts = {'metadata': 'true', 'rosbag': schema, 'pcid': pcid_str}
tabledefs.append(TableDefinition("pointcloud_formats", columns=tablecols,
options=tableopts))
for topic, infos in topics.items():
pcid = pcid_for_topic[topic]
columns, _, _, _, _, _ = get_columns(bag, topic, infos, pcid,
patch_column, patch_columns)
tablecols = [get_column_def(k, *v) for k, v in columns.items()]
tableopts = {'topic': topic, 'rosbag': schema, 'pcid': str(pcid)}
tabledefs.append(TableDefinition(topic, columns=tablecols, options=tableopts))
return tabledefs
def execute(self, quals, columns):
if self.pointcloud_formats is not None:
for f in self.pointcloud_formats:
yield f
return
self.patch_data = ''
from rospy.rostime import Time
tmin = None
tmax = None
for qual in quals:
if qual.field_name == "time":
t = int(qual.value)
t = Time(t / 1000000000, t % 1000000000)
if qual.operator in ['=', '>', '>=']:
tmin = t
if qual.operator in ['=', '<', '<=']:
tmax = t
for topic, msg, t in self.bag.read_messages(
topics=self.topic, start_time=tmin, end_time=tmax):
for row in self.get_rows(topic, msg, t, columns):
yield row
# flush leftover patch data
if self.patch_data and self.last_row:
count = int((len(self.patch_data) / self.point_size))
# in replicating mode, a single leftover point must not be reported
if count > 1 or self.patch_step_size == self.patch_size:
res = self.last_row
if self.patch_column in columns:
res[self.patch_column] = hexlify(
pack('=b3I', self.endianness, self.pcid, 0, count) + self.patch_data)
if self.patch_ply_header and 'ply' in columns:
self.ply_info['count'] = count
res['ply'] = self.patch_ply_header.format(**self.ply_info) + self.patch_data
yield res
def get_rows(self, topic, msg, t, columns, toplevel=True):
if toplevel and len(msg.__slots__) == 1:
attr = getattr(msg, msg.__slots__[0])
if isinstance(attr, list):
for msg in attr:
for row in self.get_rows(topic, msg, t, columns, False):
yield row
return
res = {}
data_columns = set(columns)
if self.patch_column in columns:
data_columns = data_columns.union(self.patch_columns) - set([self.patch_column])
if "filename" in data_columns:
res["filename"] = self.filename
if "topic" in data_columns:
res["topic"] = topic
if "time" in data_columns:
res["time"] = t.to_nsec()
if self.infos.msg_type == 'sensor_msgs/PointCloud2':
self.patch_count = self.patch_count_pointcloud or (msg.width*msg.height)
self.point_size = msg.point_step
self.patch_size = self.patch_count * self.point_size
self.patch_step_size = self.patch_size
self.endianness = 0 if msg.is_bigendian else 1
data_columns = data_columns - set(['ply', self.patch_column])
self.patch_data += msg.data
data_columns = data_columns - set(res.keys())
for column in data_columns:
attr = msg
for col in column.split('.'):
if isinstance(attr, list):
attr = tuple(getattr(a, col) for a in attr)
else:
attr = getattr(attr, col)
if hasattr(attr, "to_nsec"):
attr = attr.to_nsec()
elif hasattr(attr, "x"):
if hasattr(attr, "w"):
attr = (attr.x, attr.y, attr.z, attr.w)
else:
attr = (attr.x, attr.y, attr.z)
elif isinstance(attr, str):
fmt = self.columns[column][3]
if fmt:
attr = unpack(fmt, attr)
res[column] = attr
if self.patch_column in columns and not self.infos.msg_type == 'sensor_msgs/PointCloud2':
fmt = self.columns[self.patch_column][3]
self.patch_count = self.patch_count_default
self.point_size = calcsize(fmt)
self.patch_size = self.patch_count * self.point_size
self.patch_step_size = self.patch_size - self.point_size
self.patch_data += get_point_data(res, self.patch_columns, fmt)
res = {k: v for k, v in res.items() if k not in self.patch_columns}
if not self.patch_data:
yield res
else:
# todo: ensure current res and previous res are equal if there is some leftover
# patch_data
while len(self.patch_data) >= self.patch_size:
data = self.patch_data[0:self.patch_size]
count = int(self.patch_size / self.point_size)
res[self.patch_column] = hexlify(
pack('=b3I', self.endianness, self.pcid, 0, count) + data)
if self.patch_ply_header and 'ply' in columns:
self.ply_info = {
'endianness': 'big' if self.endianness else 'little',
'filename': self.filename,
'topic': self.topic,
'count': count
}
res['ply'] = self.patch_ply_header.format(**self.ply_info) + data
self.patch_data = self.patch_data[self.patch_step_size:]
yield res
self.last_row = res
def data_collect(file_path, link_total):
def read_msg(bag, topic, parser):
# make list with same name of "topic" at "bag" file, extract data with "parser" function from msg.
return np.array([
parser(time, msg)
for (topic_, msg, time) in bag
if topic_ == topic
])
bag=Bag(file_path)
bag = [
(topic, msg, time)
for (topic, msg, time)
in bag.read_messages()
]
# bag is an iterable of (topic, msg, time)
start_time=bag[0][2].to_sec()
def start_time_zero(extract_data):
# make start time data to zero
buff=extract_data.T
buff[0]=buff[0]-start_time
return buff.T
def quat2rpy(extract_data):
# make quaternian to roll, pitch, yaw
buff=np.zeros([extract_data.shape[0], 4])
i=0
for (t, w, x,y,z) in extract_data:
buff[i,0]=t
r = R.from_quat([w, x, y, z])
buff[i,1:]=r.as_euler('zyx', degrees=True)
#print(buff[i,1:])
i=i+1
return buff
position_list=[]
des_position_list=[]
rotation_list=[]
des_rotation_list=[]
for i in range(1,link_total+1):
position = read_msg(
bag, '/mavros_' + str(i) +'/odar/pose',
parser=lambda t, msg: (t.to_sec(), msg.pose.position.x, msg.pose.position.y, msg.pose.position.z)
)
position=start_time_zero(position)
position_list.append(position)
des_position = read_msg(
bag, '/mavros_' + str(i) +'/odar/desired_pose',
parser=lambda t, msg: (t.to_sec(), msg.pose.position.x, msg.pose.position.y, msg.pose.position.z)
)
des_position = start_time_zero(des_position)
des_position_list.append(des_position)
rotation = read_msg(
bag, '/mavros_' + str(i) +'/odar/pose',
parser=lambda t, msg: (t.to_sec(), msg.pose.orientation.w, msg.pose.orientation.x, msg.pose.position.y, msg.pose.position.z)
)
rotation = quat2rpy(start_time_zero(rotation))
rotation_list.append(rotation)
des_rotation = read_msg(
bag, '/mavros_' + str(i) +'/odar/desired_pose',
parser=lambda t, msg: (t.to_sec(), msg.pose.orientation.w, msg.pose.orientation.x, msg.pose.position.y, msg.pose.position.z)
)
des_rotation = quat2rpy(start_time_zero(des_rotation))
des_rotation_list.append(des_rotation)
return position_list, des_position_list, rotation_list, des_rotation_list
def parse_bag(bag_file_path):
bagfile = Bag(bag_file_path)
msgs_topics = [
"/drone/pose", "/gazebo/ground_truth", "/mavros/local_position/odom",
"/mavros/px4flow/ground_distance"
]
msgs_drone_pose = {
"y": [],
"z": [],
"x": [],
"roll": [],
"pitch": [],
"yaw": [],
"time": [],
"label": "EKF-SLAM"
}
msgs_ground_truth = {
"y": [],
"z": [],
"x": [],
"roll": [],
"pitch": [],
"yaw": [],
"time": [],
"label": "Ground Truth"
}
msgs_mavros_odom = {
"y": [],
"z": [],
"x": [],
"roll": [],
"pitch": [],
"yaw": [],
"time": [],
"label": "MAVROS"
}
msgs_mavros_laser = {"range": [], "time": [], "label": "Laser Range"}
for topic, message, time in bagfile.read_messages(msgs_topics):
if topic == msgs_topics[0]:
msgs_drone_pose["x"].append(message.pose.pose.position.x)
msgs_drone_pose["y"].append(message.pose.pose.position.y)
msgs_drone_pose["z"].append(message.pose.pose.position.z)
quat = message.pose.pose.orientation
euler = quat2euler([quat.x, quat.y, quat.z, quat.w])
msgs_drone_pose["roll"].append(euler[0])
msgs_drone_pose["pitch"].append(euler[1])
msgs_drone_pose["yaw"].append(euler[2])
msgs_drone_pose["time"].append(time.to_sec())
if topic == msgs_topics[1]:
msgs_ground_truth["x"].append(message.pose.pose.position.x)
msgs_ground_truth["y"].append(message.pose.pose.position.y)
msgs_ground_truth["z"].append(message.pose.pose.position.z)
quat = message.pose.pose.orientation
euler = quat2euler([quat.x, quat.y, quat.z, quat.w])
msgs_ground_truth["roll"].append(euler[0])
msgs_ground_truth["pitch"].append(euler[1])
msgs_ground_truth["yaw"].append(euler[2])
msgs_ground_truth["time"].append(time.to_sec())
if topic == msgs_topics[2]:
msgs_mavros_odom["x"].append(message.pose.pose.position.x - 3)
msgs_mavros_odom["y"].append(message.pose.pose.position.y)
msgs_mavros_odom["z"].append(message.pose.pose.position.z)
quat = message.pose.pose.orientation
euler = quat2euler([quat.x, quat.y, quat.z, quat.w])
msgs_mavros_odom["roll"].append(euler[0])
msgs_mavros_odom["pitch"].append(euler[1])
msgs_mavros_odom["yaw"].append(euler[2])
msgs_mavros_odom["time"].append(time.to_sec())
if topic == msgs_topics[3]:
msgs_mavros_laser["range"].append(message.range)
msgs_mavros_laser["time"].append(time.to_sec())
bagfile.close()
return {
"ground_truth": msgs_ground_truth,
"mavros_odom": msgs_mavros_odom,
"mavros_laser": msgs_mavros_laser,
"ekf2": msgs_drone_pose
}
def create_ground_truth(path,
time_offset=0,
bagfile='.logs/*.bag',
topic='/StereoTUM/estimated_transform',
file=''):
from rosbag import Bag
from geometry_msgs.msg import TransformStamped
from glob import glob as findfiles
if not file: file = p.join(path, 'data', 'ground_truth.csv')
first = None
msgs = []
timefiles = findfiles(p.join(path, 'params', 'time.yaml'))
if not timefiles:
raise ValueError('Can\'t find %s/params/time.yaml' % path)
return
seq_stamps = (None, None)
with open(timefiles[0]) as stream:
time = yaml.load(stream)
seq_stamps = (time['time']['start'], time['time']['end'])
print('Using time offset: %fs' % time_offset)
a = p.join(path, bagfile)
b = findfiles(a)
if not b:
raise ValueError('Can\'t find the bag file %s' % a)
return
bag = Bag(b[0])
print('Found ground truth file at: %s' % b[0])
print('Will listen to topic %s' % topic)
for bagtopic, msg, t in bag.read_messages():
if bagtopic != topic: continue
now = t.to_sec() + time_offset
if now < seq_stamps[0]: continue # before recording started
if first is None: first = now
if now > seq_stamps[1]: continue # after recording finished
stamp = now - first
# Check if any message with the same timestamp has already been added...
if any(item.get('time', None) == stamp for item in msgs): continue
# If not, add the message to the list of ground truths
msgs.append({
'time': stamp,
'trans': msg.transform.translation,
'orien': msg.transform.rotation
})
bag.close()
print('Noticed %d messages, will now sort them by time...' % len(msgs))
msgs = sorted(msgs, key=lambda x: x['time'])
with open(file, 'w') as f:
print('Creating new file %s' % file)
f.write(
'Timestamp [s]\tTranslation X [m]\tTranslation Y [m]\tTranslation Z [m]\tOrientation W\tOrientation X\tOrientation Y\tOrientation Z\n'
)
for msg in msgs:
f.write('%.6f\t' % msg['time'])
f.write('%6.4f\t' % msg['trans'].x)
f.write('%6.4f\t' % msg['trans'].y)
f.write('%6.4f\t' % msg['trans'].z)
f.write('%6.4f\t' % msg['orien'].w)
f.write('%6.4f\t' % msg['orien'].x)
f.write('%6.4f\t' % msg['orien'].y)
f.write('%6.4f' % msg['orien'].z)
f.write('\n')
# remove last newline from file
f.seek(-1, os.SEEK_END)
f.truncate()
