-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathsave_data.py
113 lines (92 loc) · 3.4 KB
/
save_data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
# Author: Srikanth Malla
# Date: 11 July 2020
# extract the sensor data from tf records and save in folders
import os
from utils import *
import glob
import tensorflow.compat.v1 as tf
from joblib import Parallel, delayed
import multiprocessing
from progressbar import ProgressBar
tf.enable_eager_execution()
## hyperparams
output_folder = "/home/smalla/waymo_data/outputs/"
# input_folder = "/home/smalla/waymo_data/v_1_2/Training/tfrecords/"
# input_folder = "/home/smalla/waymo_data/v_1_2/Validation/tfrecords/"
input_folder = "/home/smalla/waymo_data/v_1_2/Domain_Adaptation/tfrecords/"
files = glob.glob(input_folder+"*.tfrecord")
camera_images = False
pointclouds = False
top_view_lidar_image = True
camera_labels = False
lidar_labels = False
ego_vehicle_motion = False
projected_lidar_labels = False
def create_dir(folder):
if not os.path.exists(folder):
os.mkdir(folder)
def process(file):
outfolder = output_folder+file.split("/")[-1].split(".")[0]
create_dir(outfolder)
dataset = tf.data.TFRecordDataset(file, compression_type='')
## read all frames
for indx, 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)
# # 1.camera images
if camera_images:
visualize_cameras(frame, range_images, camera_projections, range_image_top_pose, outfolder, indx, save=True)
# # # 2.top view lidar image
if top_view_lidar_image:
out_file1 = outfolder+"/lidar_top"
create_dir(out_file1)
out_file = out_file1+"/"+str(indx).zfill(6)+".png"
if not os.path.isfile(out_file):
lidar_top_view(frame, range_images, camera_projections, range_image_top_pose, out_file, save=True)
# # 3.pointclouds
if pointclouds:
out_file1 = outfolder+"/pointclouds"
create_dir(out_file1)
out_file = out_file1+"/"+str(indx).zfill(6)+".npy"
if not os.path.isfile(out_file):
lidar_data(frame, range_images, camera_projections, range_image_top_pose, out_file, save=True)
#4. camera labels
if camera_labels:
labels_camera(frame, range_images, camera_projections, range_image_top_pose, outfolder, indx, save=True)
#5. pointcloud labels
if lidar_labels:
out_file1 = outfolder+"/labels_pc"
create_dir(out_file1)
out_file = out_file1+"/"+str(indx).zfill(6)+".npy"
if not os.path.isfile(out_file):
labels_pc(frame, range_images, camera_projections, range_image_top_pose, out_file, indx, save=True)
#6. ego motion
if ego_vehicle_motion:
out_file1 = outfolder+"/ego_motion"
create_dir(out_file1)
out_file = out_file1+"/"+str(indx).zfill(6)+".npy"
if not os.path.isfile(out_file):
ego_motion(frame, out_file, save=True)
#7. projection of lidar labels on to image (for the purpose of tracking ID association across sensors)
if projected_lidar_labels:
lidar_labels_on_cam(frame, outfolder, indx, save=True)
print(file)
if __name__ == '__main__':
## read all files
parallel= False
files.sort()
# batch = files[0:int(len(files)/3)]
# batch = files[int(len(files)/3):int(2*len(files)/3)]
# batch = files[int(2*len(files)/3):len(files)]
batch = files
## sequential
if not parallel:
pbar = ProgressBar()
for file in pbar(batch):
process(file)
print(file)
## parallel
if parallel:
num_cores = multiprocessing.cpu_count()
results = Parallel(n_jobs=num_cores)(delayed(process)(file) for file in files)