-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
149 lines (119 loc) · 4.43 KB
/
utils.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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import cv2
import os
import numpy as np
import matplotlib.image as mpimg
IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS = 66, 200, 3
INPUT_SHAPE = (IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS)
def load_image(data_dir, image_file):
"""
Load RGB data from a file
"""
return mpimg.imread(os.path.join(data_dir, image_file.strip()))
def crop(image):
"""
Crop the image (removing the sky at the top and the car front at the bottom)
"""
return image[60:-25, :, :]
def resize(image):
"""
Resize the image to the input shape used by the network model
"""
return cv2.resize(image, (IMAGE_WIDTH, IMAGE_HEIGHT), cv2.INTER_AREA)
def rgb2yuv(image):
"""
Convert the image from RGB to YUV (This is what the NVIDIA model does)
"""
return cv2.cvtColor(image, cv2.COLOR_RGB2YUV)
def preprocess(image):
"""
Combine all preprocess functions into one
"""
image = crop(image)
image = resize(image)
image = rgb2yuv(image)
return image
def choose_image(data_dir, center, left, right, steering_angle):
"""
Randomly choose an image from the center, left or right, and adjust
the steering angle.
"""
choice = np.random.choice(3)
if choice == 0:
return load_image(data_dir, left), steering_angle + 0.2
elif choice == 1:
return load_image(data_dir, right), steering_angle - 0.2
return load_image(data_dir, center), steering_angle
def random_flip(image, steering_angle):
"""
Randomly flipt the image left <-> right, and adjust the steering angle.
"""
if np.random.rand() < 0.5:
image = cv2.flip(image, 1)
steering_angle = -steering_angle
return image, steering_angle
def random_translate(image, steering_angle, range_x, range_y):
"""
Randomly shift the image virtially and horizontally (translation).
"""
trans_x = range_x * (np.random.rand() - 0.5)
trans_y = range_y * (np.random.rand() - 0.5)
steering_angle += trans_x * 0.002
trans_m = np.float32([[1, 0, trans_x], [0, 1, trans_y]])
height, width = image.shape[:2]
image = cv2.warpAffine(image, trans_m, (width, height))
return image, steering_angle
def random_shadow(image):
"""
Generates and adds random shadow
"""
height, width = image.shape[:2]
x1, y1 = width * np.random.rand(), 0
x2, y2 = width * np.random.rand(), height
xm, ym = np.mgrid[0:height, 0:width]
mask = np.zeros_like(image[:, :, 1])
mask[(ym - y1) * (x2 - x1) - (y2 - y1) * (xm - x1) > 0] = 1
cond = mask == np.random.randint(2)
s_ratio = np.random.uniform(low=0.2, high=0.5)
hls = cv2.cvtColor(image, cv2.COLOR_RGB2HLS)
hls[:, :, 1][cond] = hls[:, :, 1][cond] * s_ratio
return cv2.cvtColor(hls, cv2.COLOR_HLS2RGB)
def random_brightness(image):
"""
Randomly adjust brightness of the image.
"""
hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
ratio = 1.0 + 0.4 * (np.random.rand() - 0.5)
hsv[:,:,2] = hsv[:,:,2] * ratio
return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
def augument(data_dir, center, left, right, steering_angle, range_x=100, range_y=10):
"""
Generate an augumented image and adjust steering angle.
(The steering angle is associated with the center image)
"""
image, steering_angle = choose_image(data_dir, center, left, right, steering_angle)
image, steering_angle = random_flip(image, steering_angle)
image, steering_angle = random_translate(image, steering_angle, range_x, range_y)
image = random_shadow(image)
image = random_brightness(image)
return image, steering_angle
def batch_generator(data_dir, image_paths, steering_angles, batch_size, is_training):
"""
Generate training image give image paths and associated steering angles
"""
images = np.empty([batch_size, IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS])
steers = np.empty(batch_size)
while True:
i = 0
for index in np.random.permutation(image_paths.shape[0]):
center, left, right = image_paths[index]
steering_angle = steering_angles[index]
if is_training and np.random.rand() < 0.6:
image, steering_angle = augument(data_dir, center, left, right, steering_angle)
else:
image = load_image(data_dir, center)
images[i] = preprocess(image)
steers[i] = steering_angle
i += 1
if i == batch_size:
break
yield images, steers