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 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 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 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 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()
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)
