-
Notifications
You must be signed in to change notification settings - Fork 0
/
Traffic.py
123 lines (93 loc) · 3.58 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
114
115
116
117
118
119
120
121
122
123
import cv2
import numpy as np
import os
import sys
import tensorflow as tf
from keras import layers
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: object
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 category in range(NUM_CATEGORIES):
category_path = os.path.join(data_dir, str(category))
if not os.path.isdir(category_path):
continue
for image_name in os.listdir(category_path):
image_path = os.path.join(category_path, image_name)
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
if image is None:
continue
image = cv2.resize(image, (IMG_WIDTH, IMG_HEIGHT))
images.append(image)
labels.append(category)
print(f"Loaded {len(images)} images.")
print(f"Loaded {len(labels)} labels.")
return images, labels
# noinspection PyUnresolvedReferences
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.
"""
model = tf.keras.Sequential([
layers.Conv2D(
32, (3, 3), activation='tanh', input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
),
layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2)),
layers.Dropout(0.25),
layers.Conv2D(
32, (3, 3), activation='tanh', input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
),
layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2)),
layers.Dropout(0.25),
layers.Flatten(),
layers.Dense(128, activation='tanh'),
layers.Dropout(0.25),
layers.Dense(NUM_CATEGORIES, activation='softmax')
])
model.compile(loss='categorical_crossentropy',
optimizer='adam',
metrics=['accuracy', 'precision', 'recall'])
return model
if __name__ == "__main__":
main()