Press "Enter" to skip to content

利用CNN进行红绿灯识别

版权声明:本文来自Loaders Blog,未经博主允许不得转载。 https://blog.csdn.net/qibin0506/article/details/87818281

 

突然想起来我还有个博客,好久不更了,就来个使用tensorflow识别红绿灯的项目吧。

 

现在深度学习大行其道,目前最火热的深度学习框架就是tensorflow了,tensorflow大幅度的减小了我们使用深度学习的成本。今天我们就利用tensorflow来训练一个可以识别红绿灯的项目(训练数据来自MIT开源的数据集)。

 

首先先明确一下,要完成一个CNN项目的训练和使用,至少需要一下N个步骤:

 

 

    1. 收集数据(使用开源数据集、自己采集…)

 

    1. 数据整理(使用开源数据集一般可省略)

 

    1. 读取数据集

 

    1. 定义CNN模型

 

    1. 定义loss(用来计算损失)

 

    1. 定义optimizer(用来优化损失函数)

 

    1. 训练模型并保存训练结果

 

    1. 使用结果

 

 

除去那些杂七杂八的,训练过程只需要上面步骤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 数组中, imageslabels 中的数据都是一一对应的关系。最后我们使用 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
        return tf_x, tf_y, tf_y_onehot, result

 

首先我们定义了两个 placeholder 代表了我们输入的图片和标签。从上面我们得知,一个图片对应着一个label,但是对于我们的模型而言,预测的结果最好不要说的太绝对,要幺会啪啪啪打脸,所以我们使用 one_hot 函数将结果编码成一个大小为3的概率数组,对应着预测结果为红黄绿三种的概率。

 

接下来,模型的前三层是卷积操作, conv 1 使用了一个5×5的卷积核,输入的channel为图片的深度,因为是RGB图片所以这里是3,输出的channel为32,也就是会生成32个 feature map 。每个卷积操作之后都会有一个 leaky_relu 函数,作用是加入非线性化因素,也就是 激活conv 2conv 3 都是使用了3×3的卷积核。

 

上面卷积的作用其实是进行特征提取,然后的操作就是进行分类了,首先进行 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个集合进行训练。

 

使用结果

 

模型训练好了,下面我们就得使用上这些数据了,也就是进行红绿灯的预测,不过这里由于我比较懒,就没有去网上搜索图片,而是直接使用了测试数据集中的图片。

 

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%的正确率。

 

该训练结果虽然实现了对红绿灯图片的识别,但是你打开数据集可以发现图片都是切割好的,所以该结果对一个广视区域中的红绿灯识别的结果肯定不好,要实现广视区域中的红绿灯识别我们还需要在图片中进行目标检测,这些都不是本文讨论的重点了。

 

最后本文的代码都在: https://github.com/qibin0506/traffic_light_classify

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注