使用TensorFlow進行圖片分類實例

本文使用TensorFlow的基本語法（不使用高級庫），搭建神經網絡，來完成圖片分類的功能。

實現流程是：

主要分為：制作數據集、搭建CNN網絡進行訓練，使用測試集驗證、對一張實際應用的圖片進行預測。

開發環境：ancanada + python3.5 +tensorflow1.3

文件結構如圖所示：data里面存放的是圖片，net里用來存放訓練的模型，TFrecorder是制作數據集圖片預處理，resnet是搭建的resnet18層的網絡結構，之后的幾個文件見名知義。

 

1、制作數據集

制作數據集的圖片來源，可以使用Python從百度上爬去，將圖片下載到文件夾中，每一個類別分為一個文件夾，我們使用TensorFlow中的tfrecord，將圖片划為同一格式，並且為每一類圖片添加上標簽，最后生成tfrecord文件，並將其作為神經網絡的輸入。

如圖所示，bottle、paper里面存放的是瓶子和紙盒兩類圖片，record用來存放tfrecord文件，testdata是最后用來預測的圖片

 

這樣，制作數據集的代碼可以分為，生成tfrecord文件、讀取tfrecord文件、讀取的tfrecord文件的圖片分為一個batch，這三個函數進行編寫，數據集里分為訓練集和測試集，我是將每一類的圖片的前70%存儲為訓練集的tfrecord文件，剩下的30%存儲為測試集的tfrecord文件。下面是我的實現代碼：

import  tensorflow as tf
from PIL import Image 
import os
import matplotlib.pyplot as plt 
import numpy as np

path = "D:\code\resnet\data"
train_record_path = "data/record/train.tfrecords"
test_record_path = "data/record/test.tfrecords"
classes={'bottle','paper'} #人為 設定 2 類

def _byteslist(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
 
def _int64list(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def create_train_record():
    writer = tf.python_io.TFRecordWriter(train_record_path)
    NUM = 1
    for index, name in enumerate(classes):
        class_path = path + "/" + name + '/'
        l = int(len(os.listdir(class_path))*0.7)
        print("create tf "+str(index))
        for image_name in os.listdir(class_path)[:l]:
            image_path = class_path + image_name
            img = Image.open(image_path)
            img = img.resize((224, 224))
            img_raw = img.tobytes()
            example = tf.train.Example(
                    features=tf.train.Features(feature={
                            'label': _int64list(index),
                            'img_raw': _byteslist(img_raw)}))
            writer.write(example.SerializeToString())
            print('creat train record in ', NUM)
            NUM += 1
    writer.close()
    print('creat_train_record success !')
    
def create_test_record():
    writer = tf.python_io.TFRecordWriter(test_record_path)
    NUM = 1
    for index,name in enumerate(classes):
        class_path = path + "/"+name+"/"
        l = int(len(os.listdir(class_path))*0.7)
        for image_name in os.listdir(class_path)[l:]:
            image_path = class_path + image_name
            img = Image.open(image_path)
            img = img.resize((224,224))
            img_raw = img.tobytes()
            example = tf.train.Example(
                    features = tf.train.Features(feature={
                        'label':_int64list(index),
                        'img_raw':_byteslist(img_raw)}))
            writer.write(example.SerializeToString())
            print('creat test record in',NUM)
            NUM+=1
    writer.close()
    print('creat_test_record success !')
    
def read_record(filename, img_w, img_h):
    filename_queue = tf.train.string_input_producer([filename])
    reader = tf.TFRecordReader()
    _, serialize_example = reader.read(filename_queue)
    feature = tf.parse_single_example(
            serialize_example,
            features={
                    'label': tf.FixedLenFeature([], tf.int64),
                    'img_raw': tf.FixedLenFeature([], tf.string)})
    label = feature['label']
    img = feature['img_raw']
    img = tf.decode_raw(img, tf.uint8)
    img = tf.reshape(img, (224, 224, 3))
    img = tf.image.resize_image_with_crop_or_pad(img, img_w, img_h) 
    img = tf.cast(img, tf.float32)/255
    label = tf.cast(label, tf.int32)
    return img, label

def get_batch_record(filename, batch_size, img_W, img_H):
    image, label = read_record(filename, img_W, img_H)
    image_batch, label_batch= tf.train.shuffle_batch([image, label],
                                                     batch_size=batch_size,
                                                     
                                                     capacity=30,
                                                     min_after_dequeue=10)
    label_batch = tf.one_hot(label_batch,depth=2)
    return image_batch, label_batch



# if __name__ == '__main__':
#     img, label = get_batch_record(test_record_path,1,224,224)
#     print(img)

# img, label = get_batch_record(test_record_path,2, 224, 224)

# with tf.Session() as sess:
#
#     sess.run(tf.global_variables_initializer())
#     sess.run(tf.local_variables_initializer())
#     coord = tf.train.Coordinator()
#     threads = tf.train.start_queue_runners(sess, coord)
#     for i in range(200):
#         image, l =sess.run([img, label])
#         print(image[0].shape)
        
#         #print(image[1].shape)
#         print(l[0])
        
#         plt.imshow(image[0])
#         plt.show()
#     coord.request_stop()

上述代碼的說明：生成tfrecord文件時每一類圖片的label值是定義數組class的索引號，class數組是{bottle，paper}，那么bottle文件夾中圖片的標簽值（label）是0，圖片預處理基本上就兩步，第一裁剪放縮到（224,224,3），第二步將圖片的像素值歸一化為0到1之間，標簽值在形成一個batch時候，又轉換為onehot格式。

2、搭建網絡結構，訓練模型

我采用的是resnet網絡架構中的18層結構，根據resnet論文中的結構進行堆疊，為了方便我自己的理解，沒有調用高級庫。搭建的時候主要的難點殘差結構的實現，當正常卷積通道數與捷徑的通道數相等時，兩個可直接相加，通道數不等時，對捷徑做一次1*1*通道數的卷積，之后再相加。如圖所示：

 

前向傳播搭建resnet18網絡結構代碼如下：

import tensorflow as tf
def weight_variable(shape):
    initial = tf.truncated_normal(shape,stddev=0.1)
    return tf.Variable(initial)

def bias_variable(shape):
    initial = tf.constant(0.1,shape = shape)
    return tf.Variable(initial)

def conv2d(input,filter,strides,padding="SAME"):
    return tf.nn.conv2d(input,filter,strides,padding="SAME")

def resnet18(input):
    kernel_1 = weight_variable([7,7,3,64])
    bias_1 = weight_variable([64])
    layer_1 = tf.nn.relu(conv2d(input,kernel_1,strides=[1,2,2,1]) + bias_1)
    Maxpool_1 = tf.nn.max_pool(layer_1,ksize=[1,3,3,1],strides=[1,2,2,1],padding="SAME")
    
    kernel_2 = weight_variable([3,3,64,64])
    layer_2 = tf.nn.relu(conv2d(Maxpool_1,kernel_2,strides=[1,1,1,1]))
    
    kernel_3 = weight_variable([3,3,64,64])
    layer_3 = conv2d(layer_2,kernel_3,strides=[1,1,1,1]) 
    res1 = tf.nn.relu(Maxpool_1+layer_3)
    
    kernel_4 = weight_variable([3,3,64,64])
    layer_4 = tf.nn.relu(conv2d(res1,kernel_4,strides=[1,1,1,1]))
    
    kernel_5 = weight_variable([3,3,64,64])
    layer_5 = conv2d(layer_4,kernel_5,strides=[1,1,1,1])
    res2 = tf.nn.relu(res1+layer_5)
    
    kernel_6 = weight_variable([3,3,64,128])
    layer_6 = tf.nn.relu(conv2d(res2,kernel_6,strides=[1,2,2,1]))
    
    kernel_7 = weight_variable([3,3,128,128])
    layer_7 = conv2d(layer_6,kernel_7,strides=[1,1,1,1])
    kernel_line_1 = weight_variable([1,1,64,128])
    res3 = tf.nn.relu(conv2d(res2,kernel_line_1,strides=[1,2,2,1]) + layer_7)
    
    kernel_8 = weight_variable([3,3,128,128])
    layer_8 = tf.nn.relu(conv2d(res3,kernel_8,strides=[1,1,1,1]))
    
    kernel_9 = weight_variable([3,3,128,128])
    layer_9 = conv2d(layer_8,kernel_9,strides=[1,1,1,1])
    res4 = tf.nn.relu(res3+layer_9)
    
    kernel_10 = weight_variable([3,3,128,256])
    layer_10 = tf.nn.relu(conv2d(res4,kernel_10,strides=[1,2,2,1]))
    
    kernel_11 = weight_variable([3,3,256,256])
    layer_11 = conv2d(layer_10,kernel_11,strides=[1,1,1,1])
    kernel_line_2 = weight_variable([1,1,128,256]) 
    res5 = tf.nn.relu(conv2d(res4,kernel_line_2,strides=[1,2,2,1])+layer_11)
    
    kernel_12 = weight_variable([3,3,256,256])
    layer_12 = tf.nn.relu(conv2d(res5,kernel_12,strides=[1,1,1,1]))
    
    kernel_13 = weight_variable([3,3,256,256])
    layer_13 = conv2d(layer_12,kernel_13,strides=[1,1,1,1])
    res6 = tf.nn.relu(res5+layer_13)
    
    kernel_14 = weight_variable([3,3,256,512])
    layer_14 = tf.nn.relu(conv2d(res6,kernel_14,strides=[1,2,2,1]))
    
    kernel_15 = weight_variable([3,3,512,512])
    layer_15 = conv2d(layer_14,kernel_15,strides=[1,1,1,1])
    kernel_line_3 = weight_variable([1,1,256,512])
    res7 = tf.nn.relu(conv2d(res6,kernel_line_3,strides=[1,2,2,1])+ layer_15)
    
    kernel_16 = weight_variable([3,3,512,512])
    layer_16 = tf.nn.relu(conv2d(res7,kernel_16,strides=[1,1,1,1]))
    
    kernel_17 = weight_variable([3,3,512,512])
    layer_17 = conv2d(layer_16,kernel_17,strides=[1,1,1,1])
    res8 = tf.nn.relu(layer_17+res7)
    
    avgpool = tf.nn.avg_pool(res8,ksize=[1,7,7,1],strides=[1,1,1,1],padding="VALID")
    
    line = tf.reshape(avgpool,[-1,512])
    
    fc_18 = weight_variable([512,2])
    bias_18 = bias_variable([2])
    layer_18 = tf.matmul(line,fc_18)+bias_18
    
    
    return layer_18

反向傳播，輸入圖片和標簽，使得loss最小化的代碼如下：

import tensorflow as tf
import Ipynb_importer
import matplotlib.pyplot as plt
from TFrecorder import get_batch_record
import numpy as np
from resnet import resnet18
batch_size=20
filename = "data/record/train.tfrecords"
filename_test = "data/record/test.tfrecords"
num_classes = 2
img_w = 224
img_h = 224

x = tf.placeholder(tf.float32, [None, img_w, img_h, 3])
y = tf.placeholder(tf.float32, [None, num_classes])
prediction=resnet18(x)
loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y,logits=prediction))
train_step=tf.train.AdamOptimizer(0.001).minimize(loss)


correct_prediction=tf.equal(tf.argmax(y,1),tf.argmax(prediction,1))#argmax求最大的概率
accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

image_batch, label_batch = get_batch_record(filename, batch_size, img_w, img_h)
image_batch_test, label_batch_test = get_batch_record(filename_test, batch_size, img_w, img_h)
init = tf.global_variables_initializer()

saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(init)
    saver.restore(sess, "net/my_resnet18.ckpt")
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess, coord)
    #image, label = sess.run([image_batch, label_batch])
    for i in range(302):
        image, label = sess.run([image_batch, label_batch])
        #image_test, label_test = sess.run([image_batch_test, label_batch_test])
        sess.run(train_step,feed_dict={x:image,y:label})
        l = sess.run(loss,feed_dict={x:image,y:label})
        #acc = sess.run(accuracy,feed_dict={x:image_test,y:label_test})
        if i%20==0:
            print("iter: "+str(i)+" loss "+str(l))
        saver.save(sess,"net/my_resnet18.ckpt")
    coord.request_stop()  # 7
    coord.join(threads) 

代碼的說明：里面使用了多線程讀取，最后要關閉線程。還有里面加了斷點續訓，在開始訓練之前會先載入模型參數，第一次訓練如果報錯，將32行注釋就可以了。

3、測試集驗證

這個里面就沒有什么難點了，直接放代碼：

import tensorflow as tf
import Ipynb_importer
import matplotlib.pyplot as plt
from TFrecorder import get_batch_record
import numpy as np
from resnet import resnet18
batch_size=20
filename = "data/record/test.tfrecords"
num_classes = 2
img_w = 224
img_h = 224




x = tf.placeholder(tf.float32, [None, img_w, img_h, 3])
y = tf.placeholder(tf.float32, [None, num_classes])
prediction=resnet18(x)

21 #train_step=tf.train.AdamOptimizer(0.001).minimize(loss)

correct_prediction=tf.equal(tf.argmax(y,1),tf.argmax(prediction,1))#argmax求最大的概率
accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

image_batch, label_batch = get_batch_record(filename, batch_size, img_w, img_h)

init = tf.global_variables_initializer()
saver = tf.train.Saver()

with tf.Session() as sess:
    
    sess.run(init)
    module_file =  tf.train.latest_checkpoint('/my_resnet18.ckpt')
    saver.restore(sess, "net/my_resnet18.ckpt")
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess, coord)
    
    for i in range(21):
        image, label = sess.run([image_batch, label_batch])
        acc = sess.run(accuracy,feed_dict={x:image,y:label})
        43         if i%2==0:
            print(" acc: "+str(acc))
           
    coord.request_stop()  # 7
    coord.join(threads) 

4、單張圖片預測

從文件夾中讀取一張圖片之后進行預測。

代碼如下：

 

import tensorflow as tf
import Ipynb_importer
from PIL import Image 
import matplotlib.pyplot as plt
import numpy as np
from resnet import resnet18
from skimage import io, transform
import numpy as np
from pylab import mpl


image_dir = "data/testdata/4.jpg"
MODEL_SAVE_PATH = "net/my_resnet18.ckpt"
def load_image(path):
    #傳入讀入圖片的參數路徑
    img = io.imread(path) 
    #將像素歸一化處理到[0,1]
    img = img / 255.0 
    re_img = transform.resize(img, (224, 224)) 
    plt.imshow(re_img)
    plt.show()
    img_ready = re_img.reshape((1, 224, 224, 3))
    return img_ready

def prediction_result():
    x = tf.placeholder(tf.float32, [None, 224, 224, 3])
    image = load_image(image_dir)
    y=resnet18(x)
    # 利用softmax來獲取概率
    #數組的0號代表瓶子，1號位置是紙盒
    probabilities = tf.nn.softmax(y)
    # 獲取最大概率的標簽位置
    correct_prediction = tf.argmax(y, 1)
    init = tf.global_variables_initializer()
    saver = tf.train.Saver()
    
    with tf.Session() as sess:
        sess.run(init)
        saver.restore(sess, MODEL_SAVE_PATH)
        #image = sess.run(image)
        probabilities, label = sess.run([probabilities, correct_prediction], feed_dict={x:image})
        print(probabilities)
        if label == 0:
            print("this is a bottle")
            
        
        
    
def main():
    #print（'數組的0號代表瓶子，1號位置是紙盒'）
    prediction_result()
   

if __name__ == '__main__':
    main() 

 

 

最后添加的是一個應用場景的功能，可以換應用場景，相應的就要更換數據集的圖片，例如可以使用嵌入式硬件作為終端拍攝圖片上傳到服務器進行識別，我畢業設計的應用場景就是使用樹莓派控制攝像頭對物體進行拍照，上傳到服務端，返回識別結果，語音和顯示屏輸出，來教幼兒園的小孩，識字和認識物體。這個場景的實現以后有時間會更新。寫這篇博文的目的，一是為了總結學的圖片分類的知識，二是希望給想學習圖片分類應用做入門的參考。