利用CNN进行红绿灯识别
Posted 亓斌
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了利用CNN进行红绿灯识别相关的知识，希望对你有一定的参考价值。
突然想起来我还有个博客，好久不更了，就来个使用tensorflow识别红绿灯的项目吧。
引言
现在深度学习大行其道，目前最火热的深度学习框架就是tensorflow了，tensorflow大幅度的减小了我们使用深度学习的成本。今天我们就利用tensorflow来训练一个可以识别红绿灯的项目（训练数据来自MIT开源的数据集）。
首先先明确一下，要完成一个CNN项目的训练和使用，至少需要一下N个步骤：
- 收集数据（使用开源数据集、自己采集…）
- 数据整理（使用开源数据集一般可省略）
- 读取数据集
- 定义CNN模型
- 定义loss（用来计算损失）
- 定义optimizer（用来优化损失函数）
- 训练模型并保存训练结果
- 使用结果
除去那些杂七杂八的，训练过程只需要上面步骤4~7就可以了。下面我们就开始遵循上面的步骤来训练个红绿灯识别的玩意。
收集数据
因为咱们用的是开源的数据集，图片会在文章最后给出，所以这一步略…
数据整理
还是略…
读取数据集
这一步不能略了。这一步我们要做的工作是将上面的数据集中的数据读到程序中来，数据包括图片和标签，并且是一一对应的关系。
首先了解一下该开源数据集的组织结构，然后我们按照规律来读取数据：
该数据集分为两大类，一类是训练数据集，对应的就是training目录，里面有区分红黄绿三个目录，对应的就是红绿灯的三种状态；一类是测试数据集，对应的是test目录。这里已经给我们分好了训练和测试数据，所以我们可以省去切割数据集这一操作了。下面我们就把这些数据读取到程序里。
from sklearn.utils import shuffle
import numpy as np
import os
import cv2
import utils
def read_image(path):
image = cv2.imread(path)
image = cv2.resize(image, (utils.IMAGE_WIDTH, utils.IMAGE_HEIGHT))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
normalizer_image = image / 255.0 - 0.5
return normalizer_image
def read_traffic_light(training=True):
traffic_light_dir = "traffic_light_images/"
if training:
red = traffic_light_dir + "training/red/"
yellow = traffic_light_dir + "training/yellow/"
green = traffic_light_dir + "training/green/"
else:
red = traffic_light_dir + "test/red/"
yellow = traffic_light_dir + "test/yellow/"
green = traffic_light_dir + "test/green/"
images = []
labels = []
image_name = []
for f in os.listdir(red):
images.append(read_image(red + f))
labels.append(utils.RED)
image_name.append(f)
for f in os.listdir(yellow):
images.append(read_image(yellow + f))
labels.append(utils.YELLOW)
image_name.append(f)
for f in os.listdir(green):
images.append(read_image(green + f))
labels.append(utils.GREEN)
image_name.append(f)
return shuffle(np.array(images), np.array(labels), np.array(image_name))
一顿骚操作，其实要做的事情很简单，就是根据是否为训练模式，读取不同目录下的数据，将读取到images数组中，同时将改图片对应的标签读取到labels数组中，images和labels中的数据都是一一对应的关系。最后我们使用shuffle函数把顺序打乱，这一步就算完成了。
定义模型
轮到本文的核心之一了，要训练数据，首先得有个可以训练的模型，当然本文使用的是CNN，所以我们首先要定义一个CNN模型出来，然后才能训练。
import tensorflow as tf
import utils
def get_model(is_train=False, keep_prob=0.8, alpha=0.8):
with tf.variable_scope("model", reuse=not is_train):
tf_x = tf.placeholder(tf.float32, shape=[None, utils.IMAGE_HEIGHT, utils.IMAGE_WIDTH, utils.IMAGE_CHANNEL])
tf_y = tf.placeholder(tf.int32, shape=[None])
tf_y_onehot = tf.one_hot(tf_y, utils.IMAGE_CLASSIFY)
# conv 1
filter_1 = tf.Variable(tf.truncated_normal(shape=[5, 5, 3, 32], mean=utils.mu, stddev=utils.sigma))
bias_1 = tf.Variable(tf.constant(0.1, shape=[32]))
conv_1 = tf.nn.conv2d(tf_x, filter=filter_1, strides=[1, 2, 2, 1], padding='SAME') + bias_1
leaky_relu_1 = tf.nn.leaky_relu(conv_1, alpha=alpha)
# conv 2
filter_2 = tf.Variable(tf.truncated_normal(shape=[3, 3, 32, 48], mean=utils.mu, stddev=utils.sigma))
bias_2 = tf.Variable(tf.constant(0.1, shape=[48]))
conv_2 = tf.nn.conv2d(leaky_relu_1, filter=filter_2, strides=[1, 2, 2, 1], padding='SAME') + bias_2
leaky_relu_2 = tf.nn.leaky_relu(conv_2, alpha=alpha)
# conv 3
filter_3 = tf.Variable(tf.truncated_normal(shape=[3, 3, 48, 64], mean=utils.mu, stddev=utils.sigma))
bias_3 = tf.Variable(tf.constant(0.1, shape=[64]))
conv_3 = tf.nn.conv2d(leaky_relu_2, filter=filter_3, strides=[1, 2, 2, 1], padding='SAME') + bias_3
leaky_relu_3 = tf.nn.leaky_relu(conv_3, alpha=alpha)
dropout = tf.nn.dropout(leaky_relu_3, keep_prob=keep_prob)
# flatten
shape = dropout.get_shape().as_list()
flatten_size = shape[1] * shape[2] * shape[3]
flatten = tf.reshape(dropout, [-1, flatten_size])
# fc 1
filter_4 = tf.Variable(tf.truncated_normal(shape=[flatten.get_shape().as_list()[1], 100],
mean=utils.mu, stddev=utils.sigma))
bias_4 = tf.Variable(tf.constant(0.1, shape=[100]))
fc_1 = tf.matmul(flatten, filter_4) + bias_4
leaky_relu_4 = tf.nn.leaky_relu(fc_1, alpha=alpha)
# fc 2
filter_5 = tf.Variable(tf.truncated_normal(shape=[100, 50], mean=utils.mu, stddev=utils.sigma))
bias_5 = tf.Variable(tf.constant(0.1, shape=[50]))
fc_2 = tf.matmul(leaky_relu_4, filter_5) + bias_5
leaky_relu_5 = tf.nn.leaky_relu(fc_2, alpha=alpha)
# fc 3
filter_6 = tf.Variable(tf.truncated_normal(shape=[50, 10], mean=utils.mu, stddev=utils.sigma))
bias_6 = tf.Variable(tf.constant(0.1, shape=[10]))
fc_3 = tf.matmul(leaky_relu_5, filter_6) + bias_6
leaky_relu_6 = tf.nn.leaky_relu(fc_3, alpha=alpha)
# result
filter_7 = tf.Variable(tf.truncated_normal(shape=[10, utils.IMAGE_CLASSIFY],
mean=utils.mu, stddev=utils.sigma))
bias_7 = tf.Variable(tf.constant(0.1, shape=[utils.IMAGE_CLASSIFY]))
result = tf.matmul(leaky_relu_6, filter_7) + bias_7
last_layer = result if is_train else tf.nn.softmax(result)
return tf_x, tf_y, tf_y_onehot, last_layer
首先我们定义了两个placeholder代表了我们输入的图片和标签。从上面我们得知，一个图片对应着一个label，但是对于我们的模型而言，预测的结果最好不要说的太绝对，要么会啪啪啪打脸，所以我们使用one_hot函数将结果编码成一个大小为3的概率数组，对应着预测结果为红黄绿三种的概率。
接下来，模型的前三层是卷积操作，conv 1使用了一个5x5的卷积核，输入的channel为图片的深度，因为是RGB图片所以这里是3，输出的channel为32，也就是会生成32个feature map。每个卷积操作之后都会有一个leaky_relu函数，作用是加入非线性化因素，也就是激活。conv 2和conv 3都是使用了3x3的卷积核。
上面卷积的作用其实是进行特征提取，然后的操作就是进行分类了，首先进行flatten操作，该操作的作用是将上面卷积生成的feature map拍扁，对于每一个图片拍扁后都是一个一维向量。flatten完了，下面就是全连接操作了，这里就是简单的矩阵乘法，将上面拍扁的结果慢慢的向我们上面那个one_hot后向量的size上进行靠拢。最后一个操作没有使用激活函数，是因为接下来我们要使用的loss函数会默认对结果进行softmax激活，所以这里不需要多加一个激活函数。
定义loss
模型定义好了，训练之前，我们要先定义损失函数，损失函数代表着我们预测的结果和真实结果的差距有多大，定义loss的目的是我们需要不断减小这个loss，也就是预测和真实之前的差距来优化我们的网络。
def get_loss(result, y):
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=result, labels=y)
return tf.reduce_mean(cross_entropy)
定义loss很简单，首先使用softmax_cross_entropy_with_logits函数对上面模型返回的结果进行softmax激活，softmax的作用是将一个向量映射到(0, 1)区间内，也就代表了每个位置的概率。cross_entropy用来计算两个数据的混杂程度，混杂程度越小代表着两个数据越接近。在tensorflow中，将这两个操作合并到一个函数中。接下来的reduce_mean函数是对上面得到的结果取平均值，下面我们要优化的也是这个平均值。
定义optimizer
接下来，就需要定义优化函数了，最简单的优化函数就是梯度下降法，当然其他的优化方法也是在梯度下降法之上做的优化，这里我们使用Adam优化方法。
def get_optimizer(loss):
train_variables = tf.trainable_variables()
optimizer = tf.train.AdamOptimizer(learning_rate=utils.lr).minimize(loss, var_list=train_variables)
return optimizer
tensorflow提供了Adam优化的方法，我们只需要调用即可，learning_rate代表了初始化学习率。
训练模型
下面就开始训练模型了，这里我们训练50轮，最后可以得到一个损失率为1%的结果。
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.train.Saver()
with tf.Session() as sess:
sess.run(tf.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")
这里面我们迭代50轮，每一轮不断从数据集中读取128个集合进行训练。最后的loss为：
epoch: 38, loss: 0.023914596065878868
epoch: 39, loss: 0.0236150324344635
epoch: 40, loss: 0.023484617471694946
epoch: 41, loss: 0.020968815311789513
epoch: 42, loss: 0.020626120269298553
epoch: 43, loss: 0.025522083044052124
epoch: 44, loss: 0.017546119168400764
epoch: 45, loss: 0.01773049682378769
epoch: 46, loss: 0.01840411312878132
epoch: 47, loss: 0.027897536754608154
epoch: 48, loss: 0.019032606855034828
epoch: 49, loss: 0.00996661651879549
使用结果
模型训练好了，下面我们就得使用上这些数据了，也就是进行红绿灯的预测，不过这里由于我比较懒，就没有去网上搜索图片，而是直接使用了测试数据集中的图片。
import tensorflow as tf
import numpy as np
import model
import images
import utils
import random
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)))
if __name__ == '__main__':
test()
这里随机从测试数据集中读取200对数据，然后使用result.eval函数进行预测，最后能得到差不多99%的正确率。
True: 5ec72a64-1f86-42ae-bf72-f99afce85d16.jpg is green, result is green
True: 0b602dec-60ad-4601-ba50-30850e9b6895.jpg is green, result is green
True: 0bbaa6ca-9f78-4cad-8c22-b95648fb90be.jpg is green, result is green
True: 6a3dd9e4-65c3-4629-b900-ef8d94439f51.jpg is red, result is red
True: 4cd6a69f-2bb4-48b4-b64e-a0a741922b86.jpg is red, result is red
True: 4bfa49d6-5bcb-4fa5-8c9f-935f4e4527f2.jpg is red, result is red
True: 4c4d6c01-eae9-4f73-8bf4-b3dc537d8595.jpg is green, result is green
True: 4ea174a4-5956-4427-b618-d6ba3b528ab7.jpg is red, result is red
True: 07bfdf08-8e84-482f-b1a3-2e915447d248.jpg is red, result is red
True: 0bda4389-2f4a-45fe-8597-c6918a6d15b5.jpg is red, result is red
True: 2ac56486-8fd3-4289-930e-1225a1dfffe9.jpg is red, result is red
True: 5dd25b77-c117-4cca-85e9-9843e917e0f1.jpg is green, result is green
True: 5ded079d-6b51-4c15-835d-02a25e94d07c.jpg is green, result is green
True: 1e3c1562-4b74-4260-bf8a-8b12e5a1782f.jpg is red, result is red
True: 3f609861-dff1-4efa-a2b4-e8b2621eb936.jpg is red, result is red
accuracy is 0.995
总结
该训练结果虽然实现了对红绿灯图片的识别，但是你打开数据集可以发现图片都是切割好的，所以该结果对一个广视区域中的红绿灯识别的结果肯定不好，要实现广视区域中的红绿灯识别我们还需要在图片中进行目标检测，这些都不是本文讨论的重点了。
最后本文的代码都在：https://github.com/qibin0506/traffic_light_classify
以上是关于利用CNN进行红绿灯识别的主要内容，如果未能解决你的问题，请参考以下文章
利用pytorch CNN手写字母识别神经网络模型识别多手写字母（A-Z）