def train():
x_train, y_train, _ = images.read_traffic_light(True)
x_test, y_test, _ = images.read_traffic_light(False)
train_batches = x_train.shape[0]
x, y, one_hot, result = model.get_model(is_train=True)
loss = get_loss(result, one_hot)
optimizer = get_optimizer(loss)
saver = tf.compat.v1.train.Saver()
with tf.compat.v1.Session() as sess:
sess.run(tf.compat.v1.global_variables_initializer())
for epoch in range(utils.epochs):
for batch in range(train_batches // utils.batch_size):
start = batch * utils.batch_size
next_x = x_train[start:start + utils.batch_size]
next_y = y_train[start:start + utils.batch_size]
sess.run(optimizer, feed_dict={x: next_x, y: next_y})
loss_result = sess.run(loss, feed_dict={x: x_test, y: y_test})
print("epoch: {}, loss: {}".format(epoch, loss_result))
saver.save(sess, "./result/result.ckpt")
def test():
x_test, y_test, raw_names = images.read_traffic_light(False)
idxs = [random.randint(0, x_test.shape[0] - 1) for _ in range(200)]
pics = []
labels = []
names = []
for i in idxs:
pics.append(x_test[i])
labels.append(y_test[i])
names.append(raw_names[i])
x, _, _, result = model.get_model(is_train=False, keep_prob=1)
with tf.Session() as sess:
saver = tf.train.Saver()
saver.restore(sess, "./result/result.ckpt")
dists = result.eval(feed_dict={x: pics})
right_count = 0
for i in range(len(dists)):
print(i)
dist = dists[i]
pred_result = np.argmax(dist) == labels[i]
if pred_result:
right_count += 1
print("{}: {} is {}, result is {}".format(pred_result, names[i],
utils.get_traffic_name(labels[i]),
utils.get_traffic_name(np.argmax(dist))))
print("accuracy is {}".format(right_count / len(dists)))