def parse_range_image(frame,tfp):
"""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]}.
range_image_top_pose: range image pixel pose for top lidar.
"""
range_images = {}
range_image_top_pose = None
for laser in frame.lasers:
if len(laser.ri_return1.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test
range_image_str_tensor = tfp.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 = tfp.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()))
if len(laser.ri_return2.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test
range_image_str_tensor = tfp.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)
return range_images, range_image_top_pose
def parse_range_image_and_camera_projection(frame):
ranged_images = {}
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()))
ranged_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()))
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()))
ranged_images[laser.name].append(ri)
return ranged_images, range_image_top_pose
def parse_range_image_and_camera_projection(self, 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]}.
range_image_top_pose: range image pixel pose for top lidar.
"""
self.__range_images = {}
for laser in frame.lasers:
if len(laser.ri_return1.range_image_compressed) > 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()))
self.__range_images[laser.name] = [ri]
if laser.name == open_dataset.LaserName.TOP:
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()))
if len(laser.ri_return2.range_image_compressed) > 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()))
self.__range_images[laser.name].append(ri)
return self.__range_images, range_image_top_pose
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 _get_range_image_pose(self, lasers):
"""Fetches the per-pixel pose information for the range image."""
range_image_gbr_pose = None
for laser in lasers:
if laser.name != dataset_pb2.LaserName.TOP:
continue
pose_str = zlib.decompress(
laser.ri_return1.range_image_pose_compressed)
# Deserialize from MatrixFloat serialization.
range_image_gbr_pose = dataset_pb2.MatrixFloat()
range_image_gbr_pose.ParseFromString(pose_str)
assert range_image_gbr_pose is not None
shape = list(range_image_gbr_pose.shape.dims)
range_image_gbr_pose_tensor = np.array(
range_image_gbr_pose.data).reshape(shape)
range_image_gbr_pose_tensor_rotation = transform_utils.get_rotation_matrix(
range_image_gbr_pose_tensor[..., 0],
range_image_gbr_pose_tensor[..., 1],
range_image_gbr_pose_tensor[..., 2])
range_image_gbr_pose_tensor_translation = range_image_gbr_pose_tensor[
..., 3:]
range_image_gbr_pose_tensor = transform_utils.get_transform(
range_image_gbr_pose_tensor_rotation,
range_image_gbr_pose_tensor_translation)
assert range_image_gbr_pose_tensor.shape == (64, 2650, 4, 4)
return range_image_gbr_pose_tensor
def _parse_range_image(self, range_image):
"""Parse range_image proto and convert to MatrixFloat form."""
if range_image.range_image_compressed:
ri_str = zlib.decompress(range_image.range_image_compressed)
# Deserialize from MatrixFloat serialization.
ri = dataset_pb2.MatrixFloat()
ri.ParseFromString(ri_str)
else:
ri = range_image.range_image
if range_image.range_image_flow_compressed:
ri_str = zlib.decompress(range_image.range_image_flow_compressed)
# Deserialize from MatrixFloat serialization.
ri_flow = dataset_pb2.MatrixFloat()
ri_flow.ParseFromString(ri_str)
else:
ri_flow = None
return ri, ri_flow
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
from waymo_open_dataset.utils import range_image_utils
from waymo_open_dataset.utils import transform_utils
from waymo_open_dataset.utils import frame_utils
from waymo_open_dataset import dataset_pb2 as open_dataset
import hub
import zlib
from PIL import Image
filenames = os.listdir('/home/edward/waymo/validation')[0:1]
filenames = [
'/home/edward/waymo/validation/' + filename for filename in filenames
]
frame_count = 0
for filename in filenames:
dataset = tf.data.TFRecordDataset(filename)
# print('Image Count: {} File Count: {}/{}'.format(image_count, file_counter + 1, len(filenames)))
for data in dataset:
frame_count += 1
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
for laser in frame.lasers:
name = laser.name
ri1 = laser.ri_return1
ri2 = laser.ri_return2
dec = zlib.decompress(ri1.range_image_compressed)
range_image = open_dataset.MatrixFloat()
range_image.ParseFromString(dec)
print(name, range_image.shape.dims)
