forked from commaai/speedchallenge
-
Notifications
You must be signed in to change notification settings - Fork 1
/
extract_comma_struct2depth.py
87 lines (71 loc) · 2.41 KB
/
extract_comma_struct2depth.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
import numpy as np
import cv2
import glob
import os
import math
from progress.bar import Bar
SKIP = 2
WIDTH = 416
HEIGHT = 128
FX = 910
FY = 910
CX = 1164
CY = 874
OUTPUT_DIR = '/mnt/Moar Bulk/comma/struct2depth/'
def crop(img, fx, fy, cx, cy):
# Perform center cropping, preserving 50% vertically.
middle_perc = 0.38
left = 1 - middle_perc
half = left / 2
a = img[int(img.shape[0]*(half)):int(img.shape[0]*(1-half)), :]
cy /= (1 / middle_perc)
# Resize to match target height while preserving aspect ratio.
wdt = int((float(HEIGHT)*a.shape[1]/a.shape[0]))
x_scaling = float(wdt)/a.shape[1]
y_scaling = float(HEIGHT)/a.shape[0]
b = cv2.resize(a, (wdt, HEIGHT))
# Adjust intrinsics.
fx*=x_scaling
fy*=y_scaling
cx*=x_scaling
cy*=y_scaling
# Perform center cropping horizontally.
remain = b.shape[1] - WIDTH
cx /= (b.shape[1] / WIDTH)
c = b[:, math.ceil(remain/2):b.shape[1]-math.floor(remain/2)]
return c, fx, fy, cx, cy
def extract_segment(path):
segment_number = os.path.basename(path)
route_id = os.path.basename(os.path.dirname(path))
output_dir = OUTPUT_DIR + route_id + '_' + segment_number
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
triplet = []
ct = 1
cam = cv2.VideoCapture(path + "/video.hevc")
while(True):
ret, frame = cam.read()
if ret:
smallimg, fx_this, fy_this, cx_this, cy_this = crop(frame, FX, FY, CX, CY)
triplet.append(smallimg)
if len(triplet) == 3:
cmb = np.hstack(triplet)
cv2.imwrite(os.path.join(output_dir, str(ct).zfill(10) + '.png'), cmb)
f = open(os.path.join(output_dir, str(ct).zfill(10) + '_cam.txt'), 'w')
f.write(str(fx_this) + ',0.0,' + str(cx_this) + ',0.0,' + str(fy_this) + ',' + str(cy_this) + ',0.0,0.0,1.0')
f.close()
del triplet[0]
ct += 1
else:
break
segments = glob.glob("/mnt/Bulk Storage/commaai/comma2k19/*/*/*")
for s in Bar('Processing').iter(segments):
extract_segment(s)
fn = open(OUTPUT_DIR + '/' + 'train.txt', 'w')
for f in glob.glob(OUTPUT_DIR + '/*/*.png'):
if '-seg.png' in f or '-fseg.png' in f:
continue
folder_name = f.split('/')[-2]
img_name = f.split('/')[-1].replace('.png', '')
fn.write(folder_name + ' ' + img_name + '\n')
fn.close()