def save_lidar(self, frame, frame_num):
""" parse and save the lidar data in psd format
:param frame: open dataset frame proto
:param frame_num: the current frame number
:return:
"""
# range_images, range_image_top_pose = self.parse_range_image_and_camera_projection(
# frame)
# points, intensity = self.convert_range_image_to_point_cloud(
# frame,
# range_images,
# range_image_top_pose)
range_images, camera_projections, range_image_top_pose = frame_utils.parse_range_image_and_camera_projection(
frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose)
points_all = np.concatenate(points, axis=0)
intensity_all = np.ones((points_all.shape[0], 1))
point_cloud = np.column_stack((points_all, intensity_all))
pc_path = LIDAR_PATH + '/' + str(frame_num).zfill(
INDEX_LENGTH) + '.bin'
point_cloud.tofile(pc_path)
def extract_frame(frames_path, outdir_img, class_mapping, resize_ratio=1.0):
dataset = tf.data.TFRecordDataset(frames_path, compression_type='')
id_dict = {}
bboxes_all = {}
scores_all = {}
cls_inds_all = {}
track_ids_all = {}
if not os.path.exists(outdir_img):
os.mkdir(outdir_img)
for fidx, data in enumerate(dataset):
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
(range_images, camera_projections, range_image_top_pose) = (
frame_utils.parse_range_image_and_camera_projection(frame))
time = frame.context.stats.time_of_day
weather = frame.context.stats.weather
for i in range(len(frame.images)):
context_name = frame.context.name
ftm = frame.timestamp_micros
cam_name = frame.images[i].name
im = tf.image.decode_jpeg(
frame.images[i].image).numpy()[:, :, ::-1]
target_size = (int(im.shape[1] * resize_ratio),
int(im.shape[0] * resize_ratio))
im = cv2.resize(im, target_size)
cv2.imwrite(
outdir_img +
'/{}.{}.{}.png'.format(context_name, ftm, cam_name), im)
def save_lidar_points(frame, cur_save_path, use_two_returns=True):
range_images, camera_projections, range_image_top_pose = \
frame_utils.parse_range_image_and_camera_projection(frame)
points, cp_points, points_in_NLZ_flag, points_intensity, points_elongation = convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=(0, 1) if use_two_returns else (0, ))
# 3d points in vehicle frame.
points_all = np.concatenate(points, axis=0)
points_in_NLZ_flag = np.concatenate(points_in_NLZ_flag,
axis=0).reshape(-1, 1)
points_intensity = np.concatenate(points_intensity, axis=0).reshape(-1, 1)
points_elongation = np.concatenate(points_elongation,
axis=0).reshape(-1, 1)
num_points_of_each_lidar = [point.shape[0] for point in points]
save_points = np.concatenate(
[points_all, points_intensity, points_elongation, points_in_NLZ_flag],
axis=-1).astype(np.float32)
np.save(cur_save_path, save_points)
# print('saving to ', cur_save_path)
return num_points_of_each_lidar
def extract_frame(
frames_path): #, outdir_img, class_mapping, resize_ratio=1.0):
dataset = tf.data.TFRecordDataset(frames_path, compression_type='')
#id_dict = {}
#bboxes_all = {}
#scores_all = {}
#cls_inds_all = {}
#track_ids_all = {}
#if not os.path.exists(outdir_img):
# os.mkdir(outdir_img)
for fidx, data in enumerate(dataset):
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
(range_images, camera_projections, range_image_top_pose) = (
frame_utils.parse_range_image_and_camera_projection(frame))
final_frame_dict = {}
pose = frame.pose
final_frame_dict["pose"] = pose
for i in range(5):
final_frame_dict[str(i + 1) + "_image"] = frame.images[i]
final_frame_dict[
str(i + 1) +
"_calibration"] = frame.context.camera_calibrations[i]
with open(
"contexts/{}.{}.pickle".format(frame.context.name,
frame.timestamp_micros),
"wb") as f:
pickle.dump(final_frame_dict, f, protocol=pickle.HIGHEST_PROTOCOL)
def save_lidar_points(frame, cur_save_path):
range_images, camera_projections, range_image_top_pose = \
frame_utils.parse_range_image_and_camera_projection(frame)
points, cp_points, points_in_NLZ_flag, points_intensity, points_elongation = \
convert_range_image_to_point_cloud(frame, range_images, camera_projections, range_image_top_pose)
# 3d points in vehicle frame.
points_all = np.concatenate(points, axis=0)
points_in_NLZ_flag = np.concatenate(points_in_NLZ_flag, axis=0).reshape(-1, 1)
points_intensity = np.concatenate(points_intensity, axis=0).reshape(-1, 1)
points_elongation = np.concatenate(points_elongation, axis=0).reshape(-1, 1)
num_points_of_each_lidar = [point.shape[0] for point in points]
save_points = np.concatenate([
points_all, points_intensity, points_elongation, points_in_NLZ_flag
], axis=-1).astype(np.float32)
np.save(cur_save_path, save_points)
# print('saving to ', cur_save_path)
dir_name, f_name = os.path.split(cur_save_path)
f_id, _ = os.path.splitext(f_name)
h5_path = os.path.join(dir_name, '{}_image_info.h5'.format(f_id))
with h5py.File(h5_path, 'w') as f:
for im_id, camera_im in enumerate(frame.images):
im = tf.image.decode_jpeg(camera_im.image).numpy()
f.create_dataset("image_{}".format(im_id), data=im)
f.create_dataset("cp_points", data=np.concatenate(cp_points, axis=0))
return num_points_of_each_lidar
def save_lidar(self, frame, file_idx, frame_idx):
"""Parse and save the lidar data in psd format.
Args:
frame (:obj:`Frame`): Open dataset frame proto.
file_idx (int): Current file index.
frame_idx (int): Current frame index.
"""
(range_images, camera_projections, range_image_top_pose
) = frame_utils.parse_range_image_and_camera_projection(frame)
# range_images, camera_projections, range_image_top_pose = \
# parse_range_image_and_camera_projection(frame)
#convert_range_image_to_point_cloud
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
keep_polar_features=True)
# 3d points in vehicle frame.
points_all = np.concatenate(points,
axis=0) #combines 5 lidar data together
# declare new index list
i = [3, 4, 5, 1]
# create output
pointsxyzintensity_output = points_all[:, i]
pc_path = f'{self.point_cloud_save_dir}/' + \
f'{str(file_idx).zfill(3)}{str(frame_idx).zfill(3)}.bin'
pointsxyzintensity_output.astype(np.float32).tofile(pc_path)
def lidar(self,frame):
(range_images, camera_projections,range_image_top_pose) = \
frame_utils.parse_range_image_and_camera_projection(frame)
points, _ = frame_utils.convert_range_image_to_point_cloud(frame,range_images, \
camera_projections, \
range_image_top_pose)
points_all = np.concatenate(points, axis=0)
return points_all
def extract(i):
FILENAME = segments[i]
run = FILENAME.split('segment-')[-1].split('.')[0]
out_base_dir = sys.argv[2] + '/%s/' % run
if not os.path.exists(out_base_dir):
os.makedirs(out_base_dir)
dataset = tf.data.TFRecordDataset(FILENAME, compression_type='')
print(FILENAME)
pc, pc_c = [], []
camID2extr_v2c = {}
camID2intr = {}
all_static_pc = []
for frame_cnt, data in enumerate(dataset):
if frame_cnt % 2 != 0: continue ### Set the sampling rate here
print('frame ', frame_cnt)
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
extr_laser2v = np.array(
frame.context.laser_calibrations[0].extrinsic.transform).reshape(
4, 4)
extr_v2w = np.array(frame.pose.transform).reshape(4, 4)
if frame_cnt == 0:
for k in range(len(frame.context.camera_calibrations)):
cameraID = frame.context.camera_calibrations[k].name
extr_c2v =\
np.array(frame.context.camera_calibrations[k].extrinsic.transform).reshape(4, 4)
extr_v2c = np.linalg.inv(extr_c2v)
camID2extr_v2c[frame.images[k].name] = extr_v2c
fx = frame.context.camera_calibrations[k].intrinsic[0]
fy = frame.context.camera_calibrations[k].intrinsic[1]
cx = frame.context.camera_calibrations[k].intrinsic[2]
cy = frame.context.camera_calibrations[k].intrinsic[3]
k1 = frame.context.camera_calibrations[k].intrinsic[4]
k2 = frame.context.camera_calibrations[k].intrinsic[5]
p1 = frame.context.camera_calibrations[k].intrinsic[6]
p2 = frame.context.camera_calibrations[k].intrinsic[7]
k3 = frame.context.camera_calibrations[k].intrinsic[8]
camID2intr[frame.images[k].name] = np.array([[fx, 0, cx],
[0, fy, cy],
[0, 0, 1]])
# lidar point cloud
(range_images, camera_projections, range_image_top_pose) = \
frame_utils.parse_range_image_and_camera_projection(frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose)
points_all = np.concatenate(points, axis=0)
np.save('%s/frame_%03d.npy' % (out_base_dir, frame_cnt), points_all)
datalist.append(
os.path.abspath('%s/frame_%03d.npy' % (out_base_dir, frame_cnt)))
def extract_frame(frames_path,
outname,
outdir_img,
class_mapping,
resize_ratio=1.0):
dataset = tf.data.TFRecordDataset(frames_path, compression_type='')
id_dict = {}
bboxes_all = {}
scores_all = {}
cls_inds_all = {}
track_ids_all = {}
if not os.path.exists(outdir_img):
os.mkdir(outdir_img)
for fidx, data in enumerate(dataset):
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
(range_images, camera_projections, range_image_top_pose) = (
frame_utils.parse_range_image_and_camera_projection(frame))
time = frame.context.stats.time_of_day
weather = frame.context.stats.weather
im = tf.image.decode_jpeg(frame.images[0].image).numpy()[:, :, ::-1]
target_size = (int(im.shape[1] * resize_ratio),
int(im.shape[0] * resize_ratio))
im = cv2.resize(im, target_size)
labels = frame.camera_labels
if len(labels) == 0:
labels = frame.projected_lidar_labels
if len(labels) == 0:
break
assert labels[0].name == 1
boxes, types, ids = extract_labels(labels[0])
bboxes, cls_inds, track_ids = convert_kitti(boxes, types, ids, id_dict)
bboxes *= resize_ratio
scores = np.zeros(cls_inds.shape, dtype='f')
bboxes_all[fidx] = bboxes
scores_all[fidx] = scores
cls_inds_all[fidx] = cls_inds
track_ids_all[fidx] = track_ids
'''
im = cv2.resize(im, (im.shape[1] // 2, im.shape[0] // 2))
print(frame.camera_labels[0])
cv2.imshow("win", im)
cv2.waitKey(30)
print("Frame count", fidx)
'''
cv2.imwrite(outdir_img + '/%04d.png' % fidx, im)
if len(bboxes_all) > 0:
writeKITTI(outname, bboxes_all, scores_all, cls_inds_all,
track_ids_all, class_mapping)
def get_lidar(frames, frame_index):
(range_images, camera_projections, range_image_top_pose
) = frame_utils.parse_range_image_and_camera_projection(
frames[frame_index])
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frames[frame_index], range_images, camera_projections,
range_image_top_pose)
intensity_top_lidar_r0 = get_intensity(
range_images[open_dataset.LaserName.TOP][0])
return points, intensity_top_lidar_r0
def get_lidar(frame):
""" parse and save the lidar data in psd format
:param frame: open dataset frame proto
"""
range_images, camera_projections, range_image_top_pose = frame_utils.parse_range_image_and_camera_projection(
frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose)
points_all = np.concatenate(points, axis=0)
intensity_all = np.ones((points_all.shape[0], 1))
point_cloud = np.column_stack((points_all, intensity_all))
return point_cloud
def extract_laser(self, ndx, frame):
# Extract the range images, camera projects and top pose
range_images, camera_projections, range_image_top_pose = frame_utils.parse_range_image_and_camera_projection(frame)
frame.lasers.sort(key=lambda laser: laser.name)
# Using the function provided by Waymo OD to convert range image into to point clouds to visualize and save the data
points, cp_points = frame_utils.convert_range_image_to_point_cloud(frame, range_images, camera_projections, range_image_top_pose)
# Point cloud data
f_p = open("{}/{}_points.data".format(self.laser_images_dir, ndx), 'wb')
pickle.dump(points, f_p)
# Camera projects for each point cloud data
f_cp = open("{}/{}_cp_points.data".format(self.laser_images_dir, ndx), 'wb')
pickle.dump(cp_points, f_cp)
def extract_tfrecord(tfrecord):
dataset = tf.data.TFRecordDataset(tfrecord)
for data in dataset:
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
(range_images, camera_projections, range_image_top_pose
) = frame_utils.parse_range_image_and_camera_projection(frame)
for image in frame.images:
extract_single_image(frame.context.name, image,
frame.camera_labels,
frame.context.camera_calibrations,
frame.timestamp_micros)
def save_lidar(self, frame, file_idx, frame_idx):
"""Parse and save the lidar data in psd format.
Args:
frame (:obj:`Frame`): Open dataset frame proto.
file_idx (int): Current file index.
frame_idx (int): Current frame index.
"""
range_images, camera_projections, range_image_top_pose = \
parse_range_image_and_camera_projection(frame)
# First return
points_0, cp_points_0, intensity_0, elongation_0 = \
self.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=0
)
points_0 = np.concatenate(points_0, axis=0)
intensity_0 = np.concatenate(intensity_0, axis=0)
elongation_0 = np.concatenate(elongation_0, axis=0)
# Second return
points_1, cp_points_1, intensity_1, elongation_1 = \
self.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=1
)
points_1 = np.concatenate(points_1, axis=0)
intensity_1 = np.concatenate(intensity_1, axis=0)
elongation_1 = np.concatenate(elongation_1, axis=0)
points = np.concatenate([points_0, points_1], axis=0)
intensity = np.concatenate([intensity_0, intensity_1], axis=0)
elongation = np.concatenate([elongation_0, elongation_1], axis=0)
timestamp = frame.timestamp_micros * np.ones_like(intensity)
# concatenate x,y,z, intensity, elongation, timestamp (6-dim)
#Lidar elongation refers to the elongation of the pulse beyond its nominal width
point_cloud = np.column_stack(
(points, intensity, elongation, timestamp))
pc_path = f'{self.point_cloud_save_dir}/{self.prefix}' + \
f'{str(file_idx).zfill(3)}{str(frame_idx).zfill(3)}.bin'
point_cloud.astype(np.float32).tofile(pc_path)
def main():
ann_json_dict = {
'images': [],
'type': 'instances',
'annotations': [],
'categories': []
}
for class_id, class_name in label_type_map.items():
cls = {'supercategory': 'none', 'id': class_id, 'name': class_name}
ann_json_dict['categories'].append(cls)
num_examples = 0
i = 0
for tfrecord in tf.data.Dataset.list_files(FLAGS.data_dir + '/*'):
i += 1
examples = []
dataset = tf.data.TFRecordDataset(tfrecord)
for data in dataset:
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
(range_images, camera_projections, range_image_top_pose
) = frame_utils.parse_range_image_and_camera_projection(frame)
for image in frame.images:
tf_example = build_example(frame.context.name,
image,
frame.camera_labels,
frame.context.camera_calibrations,
ann_json_dict=ann_json_dict)
examples.append(tf_example.SerializeToString())
random.shuffle(examples)
writer = tf.python_io.TFRecordWriter(FLAGS.output_path +
'-%05d.tfrecord' % (i))
for example in examples:
writer.write(example)
num_examples += 1
writer.close()
print("processed {} tfrecords".format(i))
json_file_path = os.path.join(
os.path.dirname(FLAGS.output_path),
'json_' + os.path.basename(FLAGS.output_path) + '.json')
with open(json_file_path, 'w') as f:
json.dump(ann_json_dict, f)
print("wrote {} examples in {} tfrecords".format(num_examples, i))
def dataExtractor(self):
frame = open_dataset.Frame()
# Convert tfrecord to a list
datasetAsList = list(self.dataset.as_numpy_iterator())
self.totalFrames = len(datasetAsList)
# Convert to bytearray
for frameIdx in range(self.totalFrames):
frame.ParseFromString(datasetAsList[frameIdx])
range_images, camera_projections, range_image_top_pose = frame_utils.parse_range_image_and_camera_projection(
frame)
frame.lasers.sort(key=lambda laser: laser.name)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose)
self.d = {
"type_unknown": 0,
"type_vehicle": 0,
"type_pedestrian": 0,
"type_sign": 0,
"type_cyclist": 0,
"label_unknown": 0,
"label_level_1": 0,
"label_level_2": 0
}
self.process_frames(frame, frame.camera_labels)
for index, image in enumerate(frame.images):
self.extract_image(image)
self.extract_labels(image, frame)
self.assign_camera_labels(image.name)
# Write camera data to folder
self.write_camera_data(self.front_labels, self.front, 'FRONT')
self.write_camera_data(self.front_left_labels, self.front_left,
'FRONT_LEFT')
self.write_camera_data(self.front_right_labels, self.front_right,
'FRONT_RIGHT')
self.write_camera_data(self.side_left_labels, self.side_left,
'SIDE_LEFT')
self.write_camera_data(self.side_right_labels, self.side_right,
'SIDE_RIGHT')
# Write range data to folder
self.write_range_data(points, cp_points)
self.count += 1
def extract(self) -> None:
for data in self._dataset:
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
(
range_images,
camera_projections,
range_image_top_pose,
) = parse_range_image_and_camera_projection(frame)
points, r_points, i_points, e_points, cp_points = convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose
)
base_name = "{}.bin".format(frame.timestamp_micros)
# cartesian
base_dir = os.path.join(self._save_dir, "laser")
os.makedirs(base_dir, exist_ok=True)
filename = os.path.join(base_dir, base_name)
self._save(filename, points)
# range
base_dir = os.path.join(self._save_dir, "laser_r")
os.makedirs(base_dir, exist_ok=True)
filename = os.path.join(base_dir, base_name)
self._save(filename, r_points)
# intensity
base_dir = os.path.join(self._save_dir, "laser_i")
os.makedirs(base_dir, exist_ok=True)
filename = os.path.join(base_dir, base_name)
self._save(filename, i_points)
# elongation
base_dir = os.path.join(self._save_dir, "laser_e")
os.makedirs(base_dir, exist_ok=True)
filename = os.path.join(base_dir, base_name)
self._save(filename, e_points)
# cemera projection
base_dir = os.path.join(self._save_dir, "laser_cp")
os.makedirs(base_dir, exist_ok=True)
filename = os.path.join(base_dir, base_name)
self._save(filename, cp_points)
logger.info("Finish extracting laser")
def preprocess_lidar_gb(frame, ignore_difficult_instances):
""" Change red channel for LiDAR projection """
tf_examples = []
(range_images, camera_projections,
_) = frame_utils.parse_range_image_and_camera_projection(frame)
# Range image
range_image = range_images[LIDAR][0]
range_image_numpy = np.asarray(range_image.data)
range_image_numpy = np.reshape(range_image_numpy, range_image.shape.dims)
# TODO: Use other lidars as well (here channel 0: range)
range_image_numpy = range_image_numpy[..., 0]
# Change shape to (width,height,1) and cast to int
range_image_numpy = np.expand_dims(range_image_numpy,
axis=2).astype(np.int32)
# TODO: Round to closest integer?
max_val = tf.cast(MAX_LIDAR_MEASURES[0], tf.int32)
# Camera projection
camera_projection = camera_projections[LIDAR][0]
camera_projection_numpy = np.asarray(camera_projection.data)
camera_projection_numpy = np.reshape(camera_projection_numpy,
camera_projection.shape.dims)
# TODO: Use second projection (Here only use first projection channels 0,1,2)
camera_projection_numpy = camera_projection_numpy[..., :3]
# Concat lidar and camera projection
lidar_projection = np.concatenate(
[range_image_numpy, camera_projection_numpy], axis=2)
# Flatten rows and columns. Resulting shape (169600,4)
lidar_projections = np.reshape(
lidar_projection,
[lidar_projection.shape[0] * lidar_projection.shape[1], -1])
# for image_index in range(len(frame.images)):
for image_index in range(len(frame.images)):
tf_examples.append(
_process_image_lidargb(image_index, frame, lidar_projections,
max_val, ignore_difficult_instances))
return tf_examples
def frame_points(self, frame, projection=False):
"""
Returns Tensor of points in a vehicle coord. system for a given frame.
If projection is True => returns points projection on frame
"""
(range_images, camera_projections, range_image_top_pose
) = frame_utils.parse_range_image_and_camera_projection(frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=0)
if projection:
return tf.constant(np.concatenate(cp_points))
else:
return tf.constant(np.concatenate(points, axis=0))
def save_lidar(self, frame, file_idx, frame_idx):
range_images, camera_projections, range_image_top_pose = parse_range_image_and_camera_projection(
frame)
# First return
points_0, cp_points_0, intensity_0, elongation_0 = \
self.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=0
)
points_0 = np.concatenate(points_0, axis=0)
intensity_0 = np.concatenate(intensity_0, axis=0)
elongation_0 = np.concatenate(elongation_0, axis=0)
# Second return
points_1, cp_points_1, intensity_1, elongation_1 = \
self.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=1
)
points_1 = np.concatenate(points_1, axis=0)
intensity_1 = np.concatenate(intensity_1, axis=0)
elongation_1 = np.concatenate(elongation_1, axis=0)
points = np.concatenate([points_0, points_1], axis=0)
intensity = np.concatenate([intensity_0, intensity_1], axis=0)
elongation = np.concatenate([elongation_0, elongation_1], axis=0)
timestamp = frame.timestamp_micros * np.ones_like(intensity)
# concatenate x,y,z, intensity, elongation, timestamp (6-dim)
point_cloud = np.column_stack(
(points, intensity, elongation, timestamp))
pc_path = f'{self.point_cloud_save_dir}/{self.prefix}' + \
f'{str(file_idx).zfill(3)}{str(frame_idx).zfill(3)}.bin'
point_cloud.astype(np.float32).tofile(pc_path)
def save_pc(frame, dst):
range_images, camera_projections, range_image_top_pose = frame_utils.parse_range_image_and_camera_projection(
frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose)
points_ri2, cp_points_ri2 = frame_utils.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=1)
points = np.concatenate(points + points_ri2, axis=0)
# norm = np.array([[0, 0, 1], [-1, 0, 0], [0, -1, 0]])
# points = np.dot(points, norm)
# set the reflectance to 1.0 for every point
points = np.concatenate(
[points, np.ones((points.shape[0], 1), dtype=np.float32)], axis=1)
points = points.reshape(-1).astype(np.float32)
points.tofile(dst)
def save_point_cloud(frame, frame_idx, out_archive):
range_images, camera_projections, range_image_top_pose =\
frame_utils.parse_range_image_and_camera_projection(frame)
points, cp_points, channels = convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose)
points_ri2, cp_points_ri2, channels_ri2 = convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=1)
for i in range(5):
name = lidar_name_map[i + 1]
cloud = np.hstack((points[i], channels[i]))
write_file_np(out_archive, "lidar_%s/%04d.bin" % (name, frame_idx),
cloud)
cloud_ri2 = np.hstack((points_ri2[i], channels_ri2[i]))
write_file_np(out_archive, "lidar_%s_ri2/%04d.bin" % (name, frame_idx),
cloud_ri2)
def save_point_cloud(self,frame):
(range_images, camera_projections,range_image_top_pose) = frame_utils.parse_range_image_and_camera_projection(frame)
points = yzl_convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose)
points_ri2 = yzl_convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=1)
for i in range(len(points)):
range1_pts = points[i]
range2_pts = points_ri2[i]
range1_pts.tofile(self.pcd_path[i] + "range1/" +self.save_name+".bin")
range2_pts.tofile(self.pcd_path[i] + "range2/" + self.save_name+".bin")
def extract_waymo_data(filename, max_frames=150):
#FILENAME = '../waymodata/segment-10206293520369375008_2796_800_2816_800_with_camera_labels.tfrecord'
dataset = tf.data.TFRecordDataset(filename, compression_type='')
#EXTRACT FRAME AND RANGE DATA
framelist = []
for i, data in enumerate(dataset):
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
framelist.append(frame)
if i > max_frames: break
range_imagelist = []
camera_imagelist = []
range_image_toplist = []
for frame in framelist:
range_images, camera_images, range_image_top = frame_utils.parse_range_image_and_camera_projection(
frame)
range_imagelist.append(range_images)
camera_imagelist.append(camera_images)
range_image_toplist.append(range_image_top)
return framelist, range_imagelist, camera_imagelist, range_image_toplist
def parse_lidar_data(frame):
(range_images, camera_projections, range_image_top_pose) = frame_utils.parse_range_image_and_camera_projection(
frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose)
# points_ri2, cp_points_ri2 = frame_utils.convert_range_image_to_point_cloud(
# frame,
# range_images,
# camera_projections,
# range_image_top_pose,
# ri_index=1)
# 3d points in vehicle frame.
points_all = np.concatenate(points, axis=0)
# points_all_ri2 = np.concatenate(points_ri2, axis=0)
# camera projection corresponding to each point.
cp_points_all = np.concatenate(cp_points, axis=0)
# cp_points_all_ri2 = np.concatenate(cp_points_ri2, axis=0)
images = sorted(frame.images, key=lambda i: i.name)
cp_points_all_concat = np.concatenate([cp_points_all, points_all], axis=-1)
# cp_points_all_concat_tensor = tf.constant(cp_points_all_concat)
# The distance between lidar points and vehicle frame origin.
points_all_tensor = tf.norm(points_all, axis=-1, keepdims=True)
cp_points_all_tensor = tf.constant(cp_points_all, dtype=tf.int32)
mask = tf.equal(cp_points_all_tensor[..., 0], images[0].name)
cp_points_all_tensor = tf.cast(tf.gather_nd(cp_points_all_tensor, tf.where(mask)), dtype=tf.float32)
points_all_tensor = tf.gather_nd(points_all_tensor, tf.where(mask))
projected_points_all_from_raw_data = tf.concat(
[cp_points_all_tensor[..., 1:3], points_all_tensor], axis=-1).numpy()
#plot_points_on_image(projected_points_all_from_raw_data, images[0], rgba, point_size=5.0)
# plt.show()
return projected_points_all_from_raw_data
def main(data_folder, output_folder, multi_process_token=(0, 1)):
tf_records = os.listdir(data_folder)
tf_records = [x for x in tf_records if 'tfrecord' in x]
tf_records = sorted(tf_records)
for record_index, tf_record_name in enumerate(tf_records):
if record_index % multi_process_token[1] != multi_process_token[0]:
continue
print('starting for raw pc', record_index + 1, ' / ', len(tf_records), ' ', tf_record_name)
FILE_NAME = os.path.join(data_folder, tf_record_name)
dataset = tf.data.TFRecordDataset(FILE_NAME, compression_type='')
segment_name = tf_record_name.split('.')[0]
# if os.path.exists(os.path.join(output_folder, '{}.npz'.format(segment_name))):
# continue
frame_num = 0
pcs = dict()
for data in dataset:
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(data.numpy()))
# extract the points
(range_images, camera_projections, range_image_top_pose) = \
frame_utils.parse_range_image_and_camera_projection(frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose, ri_index=0)
points_all = np.concatenate(points, axis=0)
pcs[str(frame_num)] = points_all
frame_num += 1
if frame_num % 10 == 0:
print('Record {} / {} FNumber {:}'.format(record_index + 1, len(tf_records), frame_num))
print('{:} frames in total'.format(frame_num))
np.savez_compressed(os.path.join(output_folder, "{}.npz".format(segment_name)), **pcs)
def waymo_to_pytorch_offline(data_root='', idx_dataset_batch=-1):
'''
Converts tfrecords from waymo open data set to
(1) Images -> torch Tensor
(2) Lidar Range Image -> torch Tensor
(3) Labels -> dictionary of dictionaries
dictionaries with following keys:
'type': 1=TYPE_VEHICLE; 2=TYPE_PEDESTRIAN; 4=TYPE_CYCLIST
'x': upper left corner x !!labeling not as in original!!
'y': upper left corner y !!labeling not as in original!!
'width': width of corresponding bounding box !!labeling not as in original!!
'height': height of corresponding bbounding box !!labeling not as in original!!
(4) heat_maps from labels -> torch Tensor; image like
colab: there are a lot of unnessecary subdirs because of googlefilestream limitations
data is not stored in batches because the data comes in sequences of 20s.
'''
# allows __iter__() for tf record dataset
tf.compat.v1.enable_eager_execution()
# get config
if not data_root:
config = get_config()
data_root = config.dir.data.root
# read all entries in data root directory
tf_dirs = [tfd for tfd in os.listdir(data_root) if tfd.startswith('tf_')]
for idx_tf_dir, tf_dir in enumerate(tf_dirs):
# skip all non tfrecord files
tf_data_path = os.path.join(data_root, tf_dir)
for file in os.listdir(tf_data_path):
if not file.endswith('.tfrecord'):
continue
# dir names
save_path_labels = os.path.join(tf_data_path, 'labels')
save_path_images = os.path.join(tf_data_path, 'images')
save_path_lidar = os.path.join(tf_data_path, 'lidar')
save_path_heat_maps = os.path.join(tf_data_path, 'heat_maps')
# create save dirs if not exist
Path(save_path_labels).mkdir(exist_ok=True)
Path(save_path_images).mkdir(exist_ok=True)
Path(save_path_lidar).mkdir(exist_ok=True)
Path(save_path_heat_maps).mkdir(exist_ok=True)
dataset = tf.data.TFRecordDataset(
os.path.join(data_root, tf_dir,
file), compression_type='') # read tfrecord
# for all datasets stored in tfrecord
for idx_data, data in enumerate(dataset):
frame = open_dataset.Frame()
frame.ParseFromString(bytearray(
data.numpy())) # pass data from tfrecord to frame
# for all images of current frame
for idx_img, image in enumerate(
frame.images
): # can probably reduce this + next if to: image = frame.images[0]
# Only consider FRONT images
if image.name != 1: # if not CameraName == FRONT: skip; best lidar, rgb match
continue
### retrieve, convert and save rgb data
np_img = np.moveaxis(
tf.image.decode_jpeg(image.image).numpy(), -1, 0
) # frame -> np array with tensor like dims: channels,y,x
downsized_img_tensor = avgpool_tensor(torch.Tensor(np_img))
img_filename = 'img_%d_%d_%d_%d' % (
idx_dataset_batch, idx_tf_dir, idx_data, idx_img)
torch.save(downsized_img_tensor,
os.path.join(save_path_images, img_filename))
### retrieve, convert and save lidar data
(range_images, camera_projections, range_image_top_pose
) = frame_utils.parse_range_image_and_camera_projection(
frame)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections,
range_image_top_pose)
lidar_array = extract_lidar_array_from_point_cloud(
points, cp_points
) # lidar corresponds to image due to the mask in this function
range_tensor = lidar_array_to_image_like_tensor(
lidar_array)
downsized_range_tensor = pool_lidar_tensor(range_tensor)
lidar_filename = 'lidar_' + img_filename
torch.save(downsized_range_tensor,
os.path.join(save_path_lidar, lidar_filename))
### retrieve, convert and save labels
label_dict = {} # dict of dicts
labels_filename = 'labels_' + img_filename
for camera_labels in frame.camera_labels:
# ignore labels corresponding to other images
if camera_labels.name != image.name:
continue
# for all labels
for idx_label, label in enumerate(
camera_labels.labels):
label_dict[str(idx_label)] = {
'type':
label.type, # weird waymo labeling
'x':
int(label.box.center_x -
0.5 * label.box.length),
'y':
int(label.box.center_y -
0.5 * label.box.width),
'height':
int(label.box.width),
'width':
int(label.box.length)
}
save_dict(label_dict,
os.path.join(save_path_labels, labels_filename))
### create ground truth maps from labels and save
heat_map_filename = 'heat_map_' + img_filename
heat_map = create_ground_truth_maps(label_dict)
downsized_heat_map = maxpool_tensor(heat_map)
torch.save(
downsized_heat_map,
os.path.join(save_path_heat_maps, heat_map_filename))
want_small_dataset_for_testing = False
if idx_data == 9 and want_small_dataset_for_testing:
return 1
print(idx_data + 1, ' IMAGES PROCESSED')
def save_lidar(self, frame, file_idx, frame_idx):
""" parse and save the lidar data in psd format
:param frame: open dataset frame proto
:param file_idx: the current file number
:param frame_idx: the current frame number
:return:
"""
pc_path = self.point_cloud_save_dir + '/' + self.prefix + \
str(file_idx).zfill(3) + str(frame_idx).zfill(3) + '.bin'
if not self.check_file_exists(pc_path):
range_images, camera_projections, range_image_top_pose = parse_range_image_and_camera_projection(
frame)
points_0, cp_points_0, intensity_0 = self.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=0)
points_0 = np.concatenate(points_0, axis=0)
intensity_0 = np.concatenate(intensity_0, axis=0)
points_1, cp_points_1, intensity_1 = self.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=1)
points_1 = np.concatenate(points_1, axis=0)
intensity_1 = np.concatenate(intensity_1, axis=0)
points = np.concatenate([points_0, points_1], axis=0)
# print('points_0', points_0.shape, 'points_1', points_1.shape, 'points', points.shape)
intensity = np.concatenate([intensity_0, intensity_1], axis=0)
# points = points_1
# intensity = intensity_1
# reference frame:
# front-left-up (waymo) -> right-down-front(kitti)
# lidar frame:
# ?-?-up (waymo) -> front-right-up (kitti)
# print('bef\n', points)
# print('bef\n', points.dtype)
# points = np.transpose(points) # (n, 3) -> (3, n)
# tf = np.array([
# [0.0, -1.0, 0.0],
# [0.0, 0.0, -1.0],
# [1.0, 0.0, 0.0]
# ])
# points = np.matmul(tf, points)
# points = np.transpose(points) # (3, n) -> (n, 3)
# print('aft\n', points)
# print('aft\n', points.dtype)
# concatenate x,y,z and intensity
point_cloud = np.column_stack((points, intensity))
# print(point_cloud.shape)
# save
# note: must save as float32, otherwise loading errors
point_cloud.astype(np.float32).tofile(pc_path)
def parse_laser(frame, parse_label=True, output_dir="", vis=False):
'''
:param parse_label:
:param vis:
:return:
points_all: np.ndarray
[N, 4] list of 3d lidar points of length 3 or 5.
[x, y, z, intensity(optional), elongation(optional)]
laser labels: format https://github.com/waymo-research/waymo-open-dataset/blob/master/waymo_open_dataset/label.proto
message Label {
// Upright box, zero pitch and roll.
message Box {
optional double center_x = 1;
optional double center_y = 2;
optional double center_z = 3;
optional double length = 5;
optional double width = 4;
optional double height = 6;
optional double heading = 7;
}
optional Box box = 1;
enum Type {
TYPE_UNKNOWN = 0;
TYPE_VEHICLE = 1;
TYPE_PEDESTRIAN = 2;
TYPE_SIGN = 3;
TYPE_CYCLIST = 4;
}
optional Type type = 3;
enum DifficultyLevel {
UNKNOWN = 0;
LEVEL_1 = 1;
LEVEL_2 = 2;
}
optional DifficultyLevel detection_difficulty_level = 5;
optional DifficultyLevel tracking_difficulty_level = 6;
optional int32 num_lidar_points_in_box = 7;
}
'''
range_images, camera_projections, range_image_top_pose = frame_utils.parse_range_image_and_camera_projection(
frame)
points, cp_points = FrameParser.convert_range_image_to_point_cloud_V2(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=0)
# 3d points in vehicle frame.
points_all = np.concatenate(points, axis=0)
# camera projection corresponding to each point.
cp_points_all = np.concatenate(cp_points, axis=0)
laser_labels = frame.laser_labels
if parse_label:
if vis:
laser_box3d, laser_box3d_categorys, detection_difficulty_level, num_lidar_points_in_box \
= FrameParser.parse_laser_labels(laser_labels)
np.save(os.path.join(output_dir, "laser_points.npy"),
points_all)
np.save(os.path.join(output_dir, "laser_cp_points.npy"),
cp_points_all)
np.save(os.path.join(output_dir, "laser_box3d.npy"),
laser_box3d)
np.save(os.path.join(output_dir, "laser_box3d_categorys.npy"),
laser_box3d_categorys)
# FrameParser.show_laser_with_box(pcds=points_all, box3ds=boxes, categorys=categorys)
return True, points_all, cp_points_all, laser_labels
def decode_tf(FILENAME):
day_file = open('./day_segment_val.txt', 'a')
basename = os.path.basename(FILENAME)[:-9]
if not os.path.exists('./waymo_decode_val/' + basename):
os.makedirs('./waymo_decode_val/' + basename + '/image/')
os.makedirs('./waymo_decode_val/' + basename + '/depth/')
os.makedirs('./waymo_decode_val/' + basename + '/pose/')
dataset = tf.data.TFRecordDataset(FILENAME, compression_type='')
count = 0
#print(len(dataset))
frame = open_dataset.Frame()
for data in dataset:
frame.ParseFromString(bytearray(data.numpy()))
(range_images, camera_projections, range_image_top_pose
) = frame_utils.parse_range_image_and_camera_projection(frame)
# plt.figure(figsize=(25, 20))
gpspose_file = open(
'./waymo_decode_val/' + basename + '/pose/' + str(count).zfill(5) +
'_gpspose.txt', 'w')
# gpspose_file.write('%s %s %s %s %s
pose = np.array([frame.pose.transform])
# print (pose.shape)
gpspose_file.write('%f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f'%(pose[0,0],pose[0,1],pose[0,2],pose[0,3],pose[0,4]\
,pose[0,5],pose[0,6],pose[0,7],pose[0,8],pose[0,9],pose[0,10],pose[0,11],pose[0,12],pose[0,13],pose[0,14],pose[0,15]))
gpspose_file.close()
if count == 0 and frame.context.stats.time_of_day == 'Day':
day_file.write('%s\n' % (basename))
day_file.close()
for index, image in enumerate(frame.images):
if index == 0:
imageio.imwrite(
'./waymo_decode_val/' + basename + '/image/' +
str(count).zfill(5) + '.jpg',
tf.image.decode_jpeg(image.image))
break
frame.lasers.sort(key=lambda laser: laser.name)
#show_range_image(get_range_image(open_dataset.LaserName.TOP, 0), 1)
#show_range_image(get_range_image(open_dataset.LaserName.TOP, 1), 4)
points, cp_points = frame_utils.convert_range_image_to_point_cloud(
frame, range_images, camera_projections, range_image_top_pose)
points_ri2, cp_points_ri2 = frame_utils.convert_range_image_to_point_cloud(
frame,
range_images,
camera_projections,
range_image_top_pose,
ri_index=1)
# 3d points in vehicle frame.
points_all = np.concatenate(points, axis=0)
points_all_ri2 = np.concatenate(points_ri2, axis=0)
# camera projection corresponding to each point.
cp_points_all = np.concatenate(cp_points, axis=0)
cp_points_all_ri2 = np.concatenate(cp_points_ri2, axis=0)
images = sorted(frame.images, key=lambda i: i.name)
cp_points_all_concat = np.concatenate([cp_points_all, points_all],
axis=-1)
cp_points_all_concat_tensor = tf.constant(cp_points_all_concat)
# The distance between lidar points and vehicle frame origin.
points_all_tensor = tf.norm(points_all, axis=-1, keepdims=True)
cp_points_all_tensor = tf.constant(cp_points_all, dtype=tf.int32)
mask = tf.equal(cp_points_all_tensor[..., 0], images[0].name)
cp_points_all_tensor = tf.cast(tf.gather_nd(cp_points_all_tensor,
tf.where(mask)),
dtype=tf.float32)
points_all_tensor = tf.gather_nd(points_all_tensor, tf.where(mask))
projected_points_all_from_raw_data = tf.concat(
[cp_points_all_tensor[..., 1:3], points_all_tensor],
axis=-1).numpy()
saved_depth = save_depth_image(projected_points_all_from_raw_data,
tf.image.decode_jpeg(image.image))
imageio.imwrite(
'./waymo_decode_val/' + basename + '/depth/' +
str(count).zfill(5) + '.png', saved_depth)
#print(projected_points_all_from_raw_data.shape)
count += 1
