基于TensorFlow图像分类实现

train.py

训练集图片文件名称中存在类型即可。根据需分类类型修改# 训练集生成和# 测试集生成代码块中的读取方式。

import os
import numpy as np
import tensorflow as tf
from network import Network
from datagenerator import ImageDataGenerator
from datetime import datetime
import glob
from tensorflow.contrib.data import Iterator

learning_rate = 1e-4
num_epochs = 1  # 迭代次数
batch_size = 50
dropout_rate = 0.5
num_classes = 5  # 类别数量
display_step = 5

filewriter_path = "tmp/tensorboard_test"  # tensorboard文件路径
checkpoint_path = "tmp/checkpoints_test"  # 模型和参数路径

if not os.path.isdir(checkpoint_path):
    os.mkdir(checkpoint_path)

train_image_path = 'train/'  # 训练集数据路径
test_image_path = 'test/'  # 测试集数据路径

label_path = []
test_label = []

# 训练集生成
image_path = np.array(glob.glob(train_image_path + '*.jpg')).tolist()
for i in range(len(image_path)):
    if 'Bus' in image_path[i]:
        label_path.append(0)
    elif 'Microbus' in image_path[i]:
        label_path.append(1)
    elif 'Sedan' in image_path[i]:
        label_path.append(2)
    elif 'SUV' in image_path[i]:
        label_path.append(3)
    elif 'Truck' in image_path[i]:
        label_path.append(4)


# 测试集生成
test_image = np.array(glob.glob(test_image_path + '*.jpg')).tolist()
for i in range(len(test_image)):
    if 'Bus' in image_path[i]:
        test_label.append(0)
    elif 'Microbus' in image_path[i]:
        test_label.append(1)
    elif 'Sedan' in image_path[i]:
        test_label.append(2)
    elif 'SUV' in image_path[i]:
        test_label.append(3)
    elif 'Truck' in image_path[i]:
        test_label.append(4)

# 调用图片生成器，把训练集图片转换成三维数组
tr_data = ImageDataGenerator(
    images=image_path,
    labels=label_path,
    batch_size=batch_size,
    num_classes=num_classes)

# 调用图片生成器，把测试集图片转换成三维数组
test_data = ImageDataGenerator(
    images=test_image,
    labels=test_label,
    batch_size=batch_size,
    num_classes=num_classes,
    shuffle=False)

with tf.name_scope('input'):
    # 定义迭代器
    iterator = Iterator.from_structure(tr_data.data.output_types,
                                   tr_data.data.output_shapes)

    training_initalize=iterator.make_initializer(tr_data.data)
    testing_initalize=iterator.make_initializer(test_data.data)

    # 定义每次迭代的数据
    next_batch = iterator.get_next()

x = tf.placeholder(tf.float32, [batch_size, 224, 224, 3])
y = tf.placeholder(tf.float32, [batch_size, num_classes])
keep_prob = tf.placeholder(tf.float32)

# 图片数据通过网络处理
model = Network(x, keep_prob, num_classes)

# 执行整个网络图
score = model.fc8

with tf.name_scope('loss'):
    # 损失函数
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=score, labels=y))
    tf.summary.scalar('loss', loss)


with tf.name_scope('optimizer'):
    # 优化器
    train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)


# 定义网络精确度
with tf.name_scope("accuracy"):
    correct_pred = tf.equal(tf.argmax(score, 1), tf.argmax(y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    tf.summary.scalar('accuracy', accuracy)

# 把精确度加入到Tensorboard

merged_summary = tf.summary.merge_all()
writer = tf.summary.FileWriter(filewriter_path)
saver = tf.train.Saver()

# 定义一代的迭代次数
train_batches_per_epoch = int(np.floor(tr_data.data_size / batch_size))
test_batches_per_epoch = int(np.floor(test_data.data_size / batch_size))

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    #saver = tf.train.Saver()
    #saver.restore(sess, "./tmp/checkpoints_t18/model_epoch10.ckpt")

    # 把模型图加入Tensorboard
    writer.add_graph(sess.graph)

    print("{} 训练开始".format(datetime.now()))
    print("{} Tensorboard at --logdir {}".format(datetime.now(), filewriter_path))

    # 迭代所有次数
    for epoch in range(num_epochs):
        sess.run(training_initalize)
        print("{} 迭代{}次开始".format(datetime.now(), epoch + 1))

        #开始训练每一代
        for step in range(train_batches_per_epoch):
            img_batch, label_batch = sess.run(next_batch)
            sess.run(train_op, feed_dict={x: img_batch, y: label_batch, keep_prob: dropout_rate})
            if step % display_step == 0:
                s = sess.run(merged_summary, feed_dict={x: img_batch,
                                                        y: label_batch,
                                                        keep_prob: 1.})

                writer.add_summary(s, epoch * train_batches_per_epoch + step)

        # 测试模型精确度
        print("{} 测试精度".format(datetime.now()))
        sess.run(testing_initalize)
        test_acc = 0.
        test_count = 0

        for _ in range(test_batches_per_epoch):
            img_batch, label_batch = sess.run(next_batch)
            acc = sess.run(accuracy, feed_dict={x: img_batch,
                                                y: label_batch,
                                                keep_prob: 1.0})
            test_acc += acc
            test_count += 1

        test_acc /= test_count

        print("{} 精度 = {:.4f}".format(datetime.now(), test_acc))

        # 把训练好的模型存储起来
        print("{} 保存模型".format(datetime.now()))

        checkpoint_name = os.path.join(checkpoint_path, 'model_epoch' + str(epoch + 1) + '.ckpt')
        save_path = saver.save(sess, checkpoint_name)

        print("{} 迭代{}次结束".format(datetime.now(), epoch + 1), save_path)

network.py

import tensorflow as tf
import numpy as np


class Network(object):

    def __init__(self, x, keep_prob, num_classes):

        self.X = x
        self.NUM_CLASSES = num_classes
        self.KEEP_PROB = keep_prob

        self.create()

    def create(self):
        #卷积层1
        conv1_1 = conv(self.X, 9, 9, 128, 4, 4, padding='VALID', name='conv1_1')
        pool1 = max_pool(conv1_1, 2, 2, 2, 2, padding='SAME', name='pool1')
        # 卷积层2
        conv2_1 = conv(pool1, 4, 4, 256, 1, 1, padding='VALID', name='conv2_1')
        pool2 = max_pool(conv2_1, 2, 2, 2, 2, padding='SAME', name='pool2')
        # 卷积层3
        conv3_1 = conv(pool2, 3, 3, 512, 1, 1, padding='SAME', name='conv3_1')
        conv3_2 = conv(conv3_1, 3, 3, 512, 1, 1, padding='SAME', name='conv3_2')
        pool3 = max_pool(conv3_2, 2, 2, 2, 2, padding='SAME', name='pool3')
        # 卷积层4
        conv4_1 = conv(pool3, 3, 3, 256, 1, 1, padding='SAME', name='conv4_1')
        #pool4 = max_pool(conv4_2, 2, 2, 2, 2, padding='SAME', name='pool4')

        flattened = tf.reshape(conv4_1, [-1, 6*6*256])
        # 全链接6
        fc6 = fc(flattened, 6*6*256, 4096, name='fc6')
        dropout6 = dropout(fc6, self.KEEP_PROB)
        # 全链接7
        fc7 = fc(dropout6, 4096, 4096, name='fc7')
        dropout7 = dropout(fc7, self.KEEP_PROB)
        # 全链接8
        self.fc8 = fc(dropout7, 4096, self.NUM_CLASSES, name='fc8', relu=False)


def conv(x, filter_height, filter_width, num_filters, stride_y, stride_x, name,
         padding='SAME'):
    input_channels = int(x.get_shape()[-1])
    convolve = lambda i, k: tf.nn.conv2d(i, k,
                                         strides=[1, stride_y, stride_x, 1],
                                         padding=padding)

    with tf.variable_scope(name) as scope:
        weights = tf.get_variable('weights', shape=[filter_height,
                                                    filter_width,
                                                    input_channels,
                                                    num_filters])
        biases = tf.get_variable('biases', shape=[num_filters])

        conv = convolve(x, weights)
        bias = tf.reshape(tf.nn.bias_add(conv, biases), tf.shape(conv))
        relu = tf.nn.relu(bias, name=scope.name)

        return relu


def fc(x, num_in, num_out, name, relu=True):

    with tf.variable_scope(name) as scope:
        weights = tf.get_variable('weights', shape=[num_in, num_out],
                                  trainable=True)
        biases = tf.get_variable('biases', [num_out], trainable=True)
        act = tf.nn.xw_plus_b(x, weights, biases, name=scope.name)

        if relu:
            relu = tf.nn.relu(act)
            return relu
        else:
            return act


def max_pool(x, filter_height, filter_width, stride_y, stride_x, name,
             padding='SAME'):
    return tf.nn.max_pool(x, ksize=[1, filter_height, filter_width, 1],
                          strides=[1, stride_y, stride_x, 1],
                          padding=padding, name=name)


def lrn(x, radius, alpha, beta, name, bias=1.0):
    return tf.nn.local_response_normalization(x, depth_radius=radius,
                                              alpha=alpha, beta=beta,
                                              bias=bias, name=name)


def dropout(x, keep_prob):
    return tf.nn.dropout(x, keep_prob)

datagenerator.py

import tensorflow as tf
import numpy as np

from tensorflow.python.framework import dtypes
from tensorflow.python.framework.ops import convert_to_tensor
from tensorflow.contrib.data import Dataset

VGG_MEAN = tf.constant([123.68, 116.779, 103.939], dtype=tf.float32)


# 把图片数据转化为三维矩阵
class ImageDataGenerator(object):
    def __init__(self, images, labels, batch_size, num_classes, shuffle=True):

        self.img_paths = images
        self.labels = labels
        self.num_classes = num_classes
        self.data_size = len(self.labels)
        self.pointer = 0

        if shuffle:
            self._shuffle_lists()

        self.img_paths = convert_to_tensor(self.img_paths, dtype=dtypes.string)
        self.labels = convert_to_tensor(self.labels, dtype=dtypes.int32)
        data = Dataset.from_tensor_slices((self.img_paths, self.labels))
        data = data.map(self._parse_function_train, num_threads=8,
                        output_buffer_size=100 * batch_size)

        data = data.batch(batch_size)

        self.data = data

    # 打乱图片顺序
    def _shuffle_lists(self):
        path = self.img_paths
        labels = self.labels
        permutation = np.random.permutation(self.data_size)
        self.img_paths = []
        self.labels = []
        for i in permutation:
            self.img_paths.append(path[i])
            self.labels.append(labels[i])

    # 把图片生成三维数组，以及把标签转化为向量
    def _parse_function_train(self, filename, label):
        one_hot = tf.one_hot(label, self.num_classes)
        img_string = tf.read_file(filename)
        img_decoded = tf.image.decode_png(img_string, channels=3)
        img_resized = tf.image.resize_images(img_decoded, [224, 224])
        img_centered = tf.subtract(img_resized, VGG_MEAN)
        img_bgr = img_centered[:, :, ::-1]
        return img_bgr, one_hot

validate_image.py

import tensorflow as tf
from network import Network
import matplotlib.pyplot as plt
import numpy as np
import glob
from tensorflow.python.framework import dtypes
from tensorflow.python.framework.ops import convert_to_tensor
from tensorflow.contrib.data import Dataset
from tensorflow.contrib.data import Iterator

VGG_MEAN = tf.constant([123.68, 116.779, 103.939], dtype=tf.float32)
class_name = ['Bus', 'Microbus', 'Sedan', 'SUV', 'Truck']
validate_image_path = 'validate/'  # 指定验证集数据路径（根据实际情况指定验证数据集的路径）


x = tf.placeholder(tf.float32, [1, 224, 224, 3])
model = Network(x, 1, 5)
score = tf.nn.softmax(model.fc8)
max = tf.arg_max(score, 1)

t_num = 0
f_num = 0
image_path = np.array(glob.glob(validate_image_path + '*.jpg')).tolist()
fo = open("false.txt", "w")

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    saver.restore(sess, "./tmp/checkpoints_t16/model_epoch7.ckpt")
    for i in range(len(image_path)):
        img_string = tf.read_file(image_path[i])
        img_decoded = tf.image.decode_png(img_string, channels=3)
        img_resized = tf.image.resize_images(img_decoded, [224, 224])
        img_resized = img_resized[:, :, ::-1]
        img_resized = np.asarray(img_resized.eval(), dtype='uint8')
        img_resized = img_resized.reshape((1, 224, 224, 3))
        prob = sess.run(max, feed_dict={x: img_resized})[0]
        t = -1
        if 'Bus' in image_path[i]:
            t = 0
        elif 'Microbus' in image_path[i]:
            t = 1
        elif 'Sedan' in image_path[i]:
            t = 2
        elif 'SUV' in image_path[i]:
            t = 3
        elif 'Truck' in image_path[i]:
            t = 4
        if t == prob:
            t_num += 1
        else:
            f_num += 1
            fo.write(image_path[i] + '_Prediction:' + str(class_name[prob]) + '\n')

print(t_num/(t_num + f_num))