-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
141 lines (94 loc) · 3 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
import numpy as np
from matplotlib import image as img
from imgaug import augmenters as iaa
import cv2
import random
DATA_DIR = './driving_log.csv'
IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS = 66, 200, 3 # 160 320 3
INPUT_SHAPE = (IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS)
def crop(image):
"""
Crop the image (removing the sky at the top and the car front at the bottom)
"""
return image[60:-25, :, :] # remove the sky and the car front
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 load_image(image_path):
"""
return the preprocessed image
"""
image = img.imread(image_path)
return image
def zoom(image):
"""
applying zoom augmenter
"""
zoom_aug = iaa.Affine(scale=(1, 1.3))
image = zoom_aug.augment_image(image)
return image
def pan(image):
"""
applying pan augmenters
"""
pan_aug = iaa.Affine(translate_percent={"x": (-0.1, 0.1), "y": (-0.1, 0.1)})
image = pan_aug.augment_image(image)
return image
def img_random_brightness(image):
"""
applying random brightness as a form of augmentation
"""
brightness = iaa.Multiply((0.2, 1.2))
image = brightness.augment_image(image)
return image
def img_random_flip(image, steering_angle):
"""
flipping images as a form of augmentation
"""
image = cv2.flip(image, 1)
steering_angle = -steering_angle
return image, steering_angle
def random_augment(image, steering_angle):
image = img.imread(image)
if np.random.rand() < 0.5:
image = pan(image)
if np.random.rand() < 0.5:
image = zoom(image)
if np.random.rand() < 0.5:
image = img_random_brightness(image)
if np.random.rand() < 0.5:
image, steering_angle = img_random_flip(image, steering_angle)
return image, steering_angle
def batch_generator(image_paths, steering_ang, batch_size, is_training):
"""
feeding batches to our model
"""
while True:
batch_img = []
batch_steering = []
for i in range(batch_size):
random_index = random.randint(0, len(image_paths)-1)
if is_training:
image, steering = random_augment(image_paths[random_index, 0], steering_ang[random_index])
else:
image = load_image(image_paths[random_index, 0])
steering = steering_ang[random_index]
im = preprocess(image)
batch_img.append(image)
batch_steering.append(steering)
yield (np.asarray(batch_img), np.asarray(batch_steering))