用自己的图片训练图像分类问题 fine-turning

神经网络采用预训练好的inception-V3，其内部权值、偏置值均为确定的值

首先在github上下载tensorflow

注意要与电脑中的tensorflow版本相匹配 可以找一下 名字为tensorflow-frankchn的包 

下载完成后打开F:\tensorflow\tensorflow-frankchn\tensorflow\examples\image_retraining这个文件夹

后续会用到retrain等代码。

准备数据集

如果有则可以使用，若没有可以在http://www.robots.ox.ac.uk/~vgg/data/上下载喜欢的数据集来进行练习

数据集放在新建文件夹data里面，data文件夹下每种图片再分类成对应的文件夹

 

如图所示的样子建立文件夹

下载模型

# coding: utf-8

# In[1]:

import tensorflow as tf
import os
import tarfile
import requests


# In[2]:

#inception模型下载地址
inception_pretrain_model_url = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz'

#模型存放地址
inception_pretrain_model_dir = "inception_model"
if not os.path.exists(inception_pretrain_model_dir):
    os.makedirs(inception_pretrain_model_dir)
    
#获取文件名，以及文件路径
filename = inception_pretrain_model_url.split('/')[-1]
filepath = os.path.join(inception_pretrain_model_dir, filename)

#下载模型
if not os.path.exists(filepath):
    print("download: ", filename)
    r = requests.get(inception_pretrain_model_url, stream=True)
    with open(filepath, 'wb') as f:
        for chunk in r.iter_content(chunk_size=1024):
            if chunk:
                f.write(chunk)
print("finish: ", filename)
#解压文件
tarfile.open(filepath, 'r:gz').extractall(inception_pretrain_model_dir)
 
#模型结构存放文件
log_dir = 'inception_log'
if not os.path.exists(log_dir):
    os.makedirs(log_dir)

#classify_image_graph_def.pb为google训练好的模型
inception_graph_def_file = os.path.join(inception_pretrain_model_dir, 'classify_image_graph_def.pb')
with tf.Session() as sess:
    #创建一个图来存放google训练好的模型
    with tf.gfile.FastGFile(inception_graph_def_file, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        tf.import_graph_def(graph_def, name='')
    #保存图的结构
    writer = tf.summary.FileWriter(log_dir, sess.graph)
    writer.close()


# In[ ]:




# In[ ]:

开始训练，打开命令提示符输入

python3  F:/tensorflow/tensorflow-frankchn/tensorflow/examples/image_retraining/retrain.py  ^ 
--bottleneck_dir bottleneck ^
--how_many_training_steps 200 ^
--model_dir F:/firsttest  ^
--output_graph output_graph.pb ^
--output_labels output_labels.txt ^
--image_dir F:/firsttest/imags
pause

第一行是下载的tensorflow中的文件，改成对应目录即可

第二行是在当前文件夹下新建一个名为bottleneck的文件夹，来存放训练过程中的一些数据

第三行是训练次数

第四行是下载的模型地址

第五行，第六行是输出的用来测试的文件

第七行是图片文件路径，修改成自己的图片文件路径即可

训练结束会得到两个文件，这是可以用来验证的文件

测试训练好的模型

# coding: utf-8


# In[1]:


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




# In[2]:


lines = tf.gfile.GFile('retrain/output_labels.txt').readlines()  #这里改为自己的输出文件路径
uid_to_human = {} 
#一行一行读取数据
for uid,line in enumerate(lines) :
    #去掉换行符
    line=line.strip('\n')
    uid_to_human[uid] = line


def id_to_string(node_id):
    if node_id not in uid_to_human:
        return ''
    return uid_to_human[node_id]




#创建一个图来存放google训练好的模型
with tf.gfile.FastGFile('retrain/output_graph.pb', 'rb') as f: #这里改为自己的输出文件路径
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    tf.import_graph_def(graph_def, name='')




with tf.Session() as sess:
    softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
    #遍历目录
    for root,dirs,files in os.walk('retrain/images/'):    #改为自己的测试集
        for file in files: 
            #载入图片
            image_data = tf.gfile.FastGFile(os.path.join(root,file), 'rb').read()
            predictions = sess.run(softmax_tensor,{'DecodeJpeg/contents:0': image_data})#图片格式是jpg格式
            predictions = np.squeeze(predictions)#把结果转为1维数据


            #打印图片路径及名称
            image_path = os.path.join(root,file)
            print(image_path)
            #显示图片
            img=Image.open(image_path)
            plt.imshow(img)
            plt.axis('off')
            plt.show()


            #排序
            top_k = predictions.argsort()[::-1]
            print(top_k)
            for node_id in top_k:     
                #获取分类名称
                human_string = id_to_string(node_id)
                #获取该分类的置信度
                score = predictions[node_id]
                print('%s (score = %.5f)' % (human_string, score))
            print()




# In[ ]:








# In[ ]:

到此就得到了一个用自己图片训练的模型

本文采用的是Google的inception-v3结构，对于差异较大的图片识别性能较好，对于相似相近的图片可能效果不是那么理想，可以尝试使用其他预训练模型。