-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtraffic.py
113 lines (89 loc) · 3.7 KB
/
traffic.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
import cv2
import numpy as np
import os
import sys
import tensorflow as tf
from sklearn.model_selection import train_test_split
EPOCHS = 10
IMG_WIDTH = 30
IMG_HEIGHT = 30
NUM_CATEGORIES = 43
TEST_SIZE = 0.4
def main():
# Check command-line arguments
if len(sys.argv) not in [2, 3]:
sys.exit("Usage: python traffic.py data_directory [model.h5]")
# Get image arrays and labels for all image files
images, labels = load_data(sys.argv[1])
# Split data into training and testing sets
labels = tf.keras.utils.to_categorical(labels)
x_train, x_test, y_train, y_test = train_test_split(
np.array(images), np.array(labels), test_size=TEST_SIZE
)
# Get a compiled neural network
model = get_model()
# Fit model on training data
model.fit(x_train, y_train, epochs=EPOCHS)
# Evaluate neural network performance
model.evaluate(x_test, y_test, verbose=2)
# Save model to file
if len(sys.argv) == 3:
filename = sys.argv[2]
model.save(filename)
print(f"Model saved to {filename}.")
def load_data(data_dir):
"""
Load image data from directory `data_dir`.
Assume `data_dir` has one directory named after each category, numbered
0 through NUM_CATEGORIES - 1. Inside each category directory will be some
number of image files.
Return tuple `(images, labels)`. `images` should be a list of all
of the images in the data directory, where each image is formatted as a
numpy ndarray with dimensions IMG_WIDTH x IMG_HEIGHT x 3. `labels` should
be a list of integer labels, representing the categories for each of the
corresponding `images`.
"""
images, labels = [], []
# For each directory in data_dir
for dir in os.listdir(data_dir):
print(f"Loading {dir} images...")
# For each file in directory
for file in os.listdir(os.path.join(data_dir, dir)):
# Load image
img = cv2.imread(os.path.join(data_dir, dir, file))
# Resize image
img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
# Add image to images list
images.append(img)
# Add label to labels list
labels.append(int(dir))
print("Loading images complete")
# Return images and labels
return images, labels
def get_model():
"""
Returns a compiled convolutional neural network model. Assume that the
`input_shape` of the first layer is `(IMG_WIDTH, IMG_HEIGHT, 3)`.
The output layer should have `NUM_CATEGORIES` units, one for each category.
"""
# Create a convolutional neural network to classify traffic sign images using Keras.
model = tf.keras.models.Sequential([
tf.keras.layers.Conv2D(32, (3, 3), activation="relu", # 32 filters, 3x3 kernel size, relu activation
input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)), # 2x2 pooling
tf.keras.layers.Conv2D(64, (3, 3), activation="relu"), # 64 filters, 3x3 kernel size, relu activation
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)), # 2x2 pooling
tf.keras.layers.Flatten(), # flatten output
tf.keras.layers.Dense(128, activation="relu"), # 128 units, relu activation
tf.keras.layers.Dropout(0.5), # dropout layer
tf.keras.layers.Dense(NUM_CATEGORIES, activation="softmax") # output layer
])
# Compile the model using Adam optimizer and categorical crossentropy loss function.
model.compile(
optimizer="adam", # adam optimizer
loss="categorical_crossentropy", # categorical crossentropy loss function
metrics=["accuracy"] # accuracy metric
)
return model
if __name__ == "__main__":
main()