def decompress(
input_file: Bag,
output_file: Bag,
camera_info: list,
compressed_camera: list,
raw_camera: list,
) -> None:
"""
Decompress a compressed image into a raw image.
Args:
input_file: the input bag to stream from
output_file: the output bag to write data to
camera_info: a list of info topics for each compressed image topic
compressed_camera: a list of compressed image topics
raw_camera: a list of raw image topics to publish
Returns:
None
"""
# get the dims for each compressed camera info topic
dims = [get_camera_dimensions(input_file, info) for info in camera_info]
# convert the list of dims to a dictionary of compressed topics to dims
dims = dict(zip(compressed_camera, dims))
# setup a dictionary mapping compressed topics to raw camera topics
raw_camera = dict(zip(compressed_camera, raw_camera))
# get the number of camera info topics to leave out
camera_info_total = input_file.get_message_count(topic_filters=camera_info)
# get the total number of output messages as total minus camera info topics
total = input_file.get_message_count() - camera_info_total
# create a progress bar for iterating over the messages in the output bag
with tqdm(total=total) as prog:
# iterate over the messages in this input bag
for topic, msg, time in input_file:
# don't write the camera info to the bag
if topic in camera_info:
continue
# if the topic is the compress image topic, decompress the image
if topic in compressed_camera:
# get the image from the compressed topic
img = get_camera_image(msg.data, dims[topic])
# create a raw image message and replace the message
msg = image_msg(img, time, dims[topic], 'rgb8')
# reset the topic
topic = raw_camera[topic]
# update the progress bar with a single iteration
prog.update(1)
# write the message
output_file.write(topic, msg, time)
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 shift(
input_file: Bag,
output_file: Bag,
shift_topic: list,
shift_sec: float,
) -> None:
"""
Shift a topic in a bag forward or backwards in time.
Args:
input_file: the input bag to stream from
output_file: the output bag to write data to
topic: the topic to shift in time
shift_sec: the amount of time to shift by
Returns:
None
"""
# create a progress bar for iterating over the messages in the bag
with tqdm(total=input_file.get_message_count()) as prog:
# iterate over the messages in this input bag
for topic, msg, time in input_file:
# if the topic is the sentinel topic, increase the count
if topic == shift_topic:
# adjust time by the shift value
time += rospy.Duration(shift_sec)
# update the message header with the new time stamp
msg.header.stamp = time
# update the progress bar with a single iteration
prog.update(1)
# write the message
output_file.write(topic, msg, time)
def stream(input_file: Bag, output_file: Bag, topics: list) -> None:
"""
Stream data from an input bag to an output bag.
Args:
input_file: the input bag to stream from
output_file: the output bag to write data to
topics: a list of the topics to include
Returns:
None
"""
included = 0
skipped = 0
# create a progress bar for iterating over the messages in the bag
with tqdm(total=input_file.get_message_count(), unit='message') as prog:
# iterate over the messages in this input bag
for topic, msg, time in input_file:
# check for matches between the topics filter and this topic
if any(fnmatchcase(topic, pattern) for pattern in topics):
# write this message to the output bag
output_file.write(topic, msg, time)
# increment the counter of included messages
included += 1
else:
# increment the counter of excluded messages
skipped += 1
# update the progress bar with a single iteration
prog.update(1)
# update the progress bar post fix text with statistics
prog.set_postfix(included=included, skipped=skipped)
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 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 _read_info(self):
"""Read some metadata from inside this rosbag."""
from rosbag import Bag, ROSBagUnindexedException, ROSBagException
try:
b = Bag(self.path)
except ROSBagUnindexedException:
b = Bag(self.path, allow_unindexed=True)
print('Reindexing', self.filename)
b.reindex()
try:
duration = b.get_end_time() - b.get_start_time()
except ROSBagException:
duration = 0
return {
'n_messages': b.get_message_count(),
'duration': duration,
'topic_list': b.get_type_and_topic_info()[1].keys()
}
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()
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)
