forked from commaai/speedchallenge
-
Notifications
You must be signed in to change notification settings - Fork 1
/
extract_combined.py
41 lines (31 loc) · 1.16 KB
/
extract_combined.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
import tensorflow as tf
import glob
import random
import extract_commaai_2k19
import extract_waymo
def write_tfrecord(dirs, out):
with tf.python_io.TFRecordWriter(out) as writer:
example_buffer = []
for i, dir in enumerate(dirs):
if "waymo" in dir:
segment = extract_waymo.extract_segment(dir)
else:
segment = extract_commaai_2k19.extract_segment(dir)
example_buffer.extend(segment)
print("extracted {} segments".format(i))
if len(example_buffer) > 1000000:
random.shuffle(example_buffer)
for e in example_buffer:
writer.write(e)
example_buffer = []
print("\n\nflushed the buffer\n\n")
random.shuffle(example_buffer)
for e in example_buffer:
writer.write(e)
waymo_segments = glob.glob("/mnt/Bulk/waymo/*/*.tfrecord")
commaai_segments = glob.glob("/mnt/Bulk/commaai/comma2k19/*/*/*")
segments = []
segments.extend(waymo_segments)
segments.extend(commaai_segments)
random.shuffle(segments)
write_tfrecord(segments, "/mnt/Bulk/monolithic_combined.tfrecord")