训练自己的分类模型
基于inception v3
介绍
下载 inception v3 模型
http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
其中的 ++classify_image_graph_def.pb++ 文件就是模型文件，但是pb文件只能使用，不能继续训练。
++imagenet_synset_to_human_label_map.txt++ ++imagenet_2012_challenge_label_map_proto.pbtxt++ 这两个文件是数字编号对应物品英文描述的索引文件
使用代码
import tensorflow as tf
import os
import numpy as np
import re
from PIL import Image
import matplotlib.pyplot as plt
class NodeLookup(object):
def __init__(self):
label_lookup_path = 'inception_model/imagenet_2012_challenge_label_map_proto.pbtxt'
uid_lookup_path = 'inception_model/imagenet_synset_to_human_label_map.txt'
self.node_lookup = self.load(label_lookup_path, uid_lookup_path)
def load(self, label_lookup_path, uid_lookup_path):
# 加载分类字符串n********对应分类名称的文件
proto_as_ascii_lines = tf.gfile.GFile(uid_lookup_path).readlines()
uid_to_human = {}
# 一行一行读取数据
for line in proto_as_ascii_lines :
# 去掉换行符
line=line.strip('\n')
# 按照'\t'分割
parsed_items = line.split('\t') # 空格
# 获取分类编号
uid = parsed_items[0]
# 获取分类名称
human_string = parsed_items[1]
# 保存编号字符串n********与分类名称映射关系
uid_to_human[uid] = human_string
# 加载分类字符串n********对应分类编号1-1000的文件
proto_as_ascii = tf.gfile.GFile(label_lookup_path).readlines()
node_id_to_uid = {}
for line in proto_as_ascii:
# 去掉换行符
line=line.strip('\n')
if line.startswith(' target_class:'):
# 获取分类编号1-1000
target_class = int(line.split(': ')[1])
if line.startswith(' target_class_string:'):
# 获取编号字符串n********
target_class_string = line.split(': ')[1]
# 保存分类编号1-1000与编号字符串n********映射关系
node_id_to_uid[target_class] = target_class_string[1:-1] # 不要冒号
# 建立分类编号1-1000对应分类名称的映射关系
node_id_to_name = {}
for key, val in node_id_to_uid.items():
# 获取分类名称
name = uid_to_human[val]
# 建立分类编号1-1000到分类名称的映射关系
node_id_to_name[key] = name
return node_id_to_name
# 传入分类编号1-1000返回分类名称
def id_to_string(self, node_id):
if node_id not in self.node_lookup:
return ''
return self.node_lookup[node_id]
# 创建一个图来存放google训练好的模型
with tf.gfile.GFile('inception_model/classify_image_graph_def.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('softmax:0')
# 遍历目录
for root,dirs,files in os.walk('images/'):
for file in files:
# 载入图片
image_data = tf.gfile.GFile(os.path.join(root,file), 'rb').read()
predictions = sess.run(softmax_tensor,{'DecodeJpeg/contents:0': image_data})#图片格式是jpg格式
predictions = np.squeeze(predictions)#把结果转为1维数据
print(predictions.shape)
# 打印图片路径及名称
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()[-5:][::-1] # 因为数据是从小到大排列的，所以取最后5个 并降序展示
node_lookup = NodeLookup()
for node_id in top_k:
# 获取分类名称
human_string = node_lookup.id_to_string(node_id)
# 获取该分类的置信度
score = predictions[node_id]
print('%s (score = %.5f)' % (human_string, score))
print()
训练自己的模型
Google官方的retrain.py
运行retrain.bat来训练自己的分类模型
python retrain.py ^
--bottleneck_dir bottleneck ^
--how_many_training_steps 200 ^
--model_dir inception_model ^
--output_graph output_graph.pb ^
--output_labels output_labels.txt ^
--image_dir data
pause切记要切换到自己的python环境再运行retrain.bat
参数解释
python retrain.py # 要运行的python文件
--bottleneck_dir bottleneck # 需要指定bottleneck的路径 可以理解为图片的特征
--how_many_training_steps 200 #训练200批次
--model_dir inception_model #下载好的inception v3模型位置
--output_graph output_graph.pb # 输出文件
--output_labels output_labels.txt # 输出标签
--image_dir data # 需要训练的图片存放路径，要用文件夹分好类，文件夹名字要有意义。
pause因为只训练全连接层，所以速度很快。
不需要大量数据集。
不需要大量批次。
属于迁移学习。
用训练好的 模型做预测
import tensorflow as tf
import os
import numpy as np
import re
from PIL import Image
import matplotlib.pyplot as plt
lines = tf.gfile.GFile('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]
# 创建一个图来存放训练好的模型
with tf.gfile.GFile('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:
# final_result为输出tensor的名字
softmax_tensor = sess.graph.get_tensor_by_name('final_result:0')
# 遍历目录
for root,dirs,files in os.walk('images/'):
for file in files:
# 载入图片
image_data = tf.gfile.GFile(os.path.join(root,file), 'rb').read()
# 把图像数据传入模型获得模型输出结果
predictions = sess.run(softmax_tensor,{'DecodeJpeg/contents:0': image_data})
# 把结果转为1维数据
predictions = np.squeeze(predictions)
# 打印图片路径及名称
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]
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()
模型准确率高是因为inception v3效果很好，而且分类难度较低。