def parse_range_image_and_camera_projection(frame):
"""Parse range images and camera projections given a frame.
Args:
frame: open dataset frame proto
Returns:
range_images: A dict of {laser_name,
[range_image_first_return, range_image_second_return]}.
camera_projections: A dict of {laser_name,
[camera_projection_from_first_return,
camera_projection_from_second_return]}.
range_image_top_pose: range image pixel pose for top lidar.
"""
range_images = {}
camera_projections = {}
range_image_top_pose = None
for laser in frame.lasers:
if len(laser.ri_return1.range_image_compressed) > 0:
# use tf.io.decode_compressed() if TF 2.0
range_image_str_tensor = tf.decode_compressed(
laser.ri_return1.range_image_compressed, 'ZLIB')
ri = open_dataset.MatrixFloat()
ri.ParseFromString(bytearray(range_image_str_tensor.numpy()))
range_images[laser.name] = [ri]
if laser.name == open_dataset.LaserName.TOP:
# use tf.io.decode_compressed() if TF 2.0
range_image_top_pose_str_tensor = tf.decode_compressed(
laser.ri_return1.range_image_pose_compressed, 'ZLIB')
range_image_top_pose = open_dataset.MatrixFloat()
range_image_top_pose.ParseFromString(
bytearray(range_image_top_pose_str_tensor.numpy()))
# use tf.io.decode_compressed() if TF 2.0
camera_projection_str_tensor = tf.decode_compressed(
laser.ri_return1.camera_projection_compressed, 'ZLIB')
cp = open_dataset.MatrixInt32()
cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy()))
camera_projections[laser.name] = [cp]
if len(laser.ri_return2.range_image_compressed) > 0:
# use tf.io.decode_compressed() if TF 2.0
range_image_str_tensor = tf.decode_compressed(
laser.ri_return2.range_image_compressed, 'ZLIB')
ri = open_dataset.MatrixFloat()
ri.ParseFromString(bytearray(range_image_str_tensor.numpy()))
range_images[laser.name].append(ri)
# use tf.io.decode_compressed() if TF 2.0
camera_projection_str_tensor = tf.decode_compressed(
laser.ri_return2.camera_projection_compressed, 'ZLIB')
cp = open_dataset.MatrixInt32()
cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy()))
camera_projections[laser.name].append(cp)
return range_images, camera_projections, range_image_top_pose
def parse_range_image_and_camera_projection(frame):
range_images = {}
camera_projections = {}
range_image_top_pose = None
for laser in frame.lasers:
if len(laser.ri_return1.range_image_compressed) > 0:
range_image_str_tensor = tf.io.decode_compressed(laser.ri_return1.range_image_compressed, 'ZLIB')
ri = open_dataset.MatrixFloat()
ri.ParseFromString(bytearray(range_image_str_tensor.numpy()))
range_images[laser.name] = [ri]
if laser.name == open_dataset.LaserName.TOP:
range_image_top_pose_str_tensor = tf.io.decode_compressed(laser.ri_return1.range_image_pose_compressed, 'ZLIB')
range_image_top_pose = open_dataset.MatrixFloat()
range_image_top_pose.ParseFromString(bytearray(range_image_top_pose_str_tensor.numpy()))
camera_projection_str_tensor = tf.io.decode_compressed(laser.ri_return1.camera_projection_compressed, 'ZLIB')
cp = open_dataset.MatrixInt32()
cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy()))
camera_projections[laser.name] = [cp]
if len(laser.ri_return2.range_image_compressed) > 0:
range_image_str_tensor = tf.io.decode_compressed(laser.ri_return2.range_image_compressed, 'ZLIB')
ri = open_dataset.MatrixFloat()
ri.ParseFromString(bytearray(range_image_str_tensor.numpy()))
range_images[laser.name].append(ri)
camera_projection_str_tensor = tf.io.decode_compressed(laser.ri_return2.camera_projection_compressed, 'ZLIB')
cp = open_dataset.MatrixInt32()
cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy()))
camera_projections[laser.name].append(cp)
return range_images, camera_projections, range_image_top_pose
def extract_record_batch(self, ds: hub.Dataset, arr: Dict[str, np.ndarray], index: int, batch) -> int:
cnt = 0
for key in ['images', 'lasers_range_image', 'lasers_camera_projection']:
arr[key] = np.zeros(ds[key].chunk, ds[key].dtype)
for i, sample in enumerate(batch):
t1 = time.time()
cnt += 1
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(sample.numpy()))
key = 'images'
arrkey = arr[key]
for image in frame.images:
if image.name > 0:
img = np.array(Image.open(io.BytesIO(bytearray(image.image))))
arrkey[i, image.name - 1, :img.shape[0], :img.shape[1]] = img
arrkey = arr['labels'][index + i]
for j, key in enumerate(['camera_labels', 'projected_lidar_labels']):
for image in getattr(frame, key):
if image.name > 0:
_count = 0
for l, label in enumerate(image.labels):
box = label.box
x = np.array([box.center_x, box.center_y, box.center_z, box.length, box.width, box.height, box.heading])
arrkey[5 * j + image.name, l] = x
_count += 1
ds['labels'].darray[index + i][5 * j + image.name] = (_count, 7)
key = 'laser_labels'
_count = 0
for j, label in enumerate(getattr(frame, key)):
box = label.box
x = np.array([box.center_x, box.center_y, box.center_z, box.length, box.width, box.height, box.heading])
arrkey[0, j] = x
_count += 1
ds['labels'].darray[index + i][0] = (_count, 7)
for kid, key in enumerate(['lasers_range_image', 'lasers_camera_projection']):
arrkey = arr[key]
for laser in frame.lasers:
if laser.name > 0:
for j, ri in enumerate([laser.ri_return1, laser.ri_return2]):
attr_name = '_'.join(key.split('_')[1:]) + '_compressed'
data = getattr(ri, attr_name)
mt = open_dataset.MatrixFloat() if kid == 0 else open_dataset.MatrixInt32()
mt.ParseFromString(zlib.decompress(data))
_arr = np.reshape(
np.array(mt.data),
tuple(mt.shape.dims),
)
arrkey[i, laser.name - 1, j, :_arr.shape[0], :_arr.shape[1]] = _arr
# print(f'{laser.name - 1}, {key}: {_arr.shape}')
t2 = time.time()
return cnt