inception_v3模型
- tensorflow_inception_graph.pb
这个文件可以自己去下载,大部分都是一样的。 - 训练需要的图片相同的类别放到同一个文件夹里
- 所有的类别图片放到一个总文件夹里
形如:
代码如下
import tensorflow as tf
import pickle
import os
import numpy as np
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
model_path = './model/tensorflow_inception_graph.pb'
input_img_path = 'C:/Users/machenike/Desktop/imgs'
output_folder = './data/bb'
if not os.path.exists(output_folder):
os.mkdir(output_folder)
def parse_img(input_img_path):
img_datas = []
img_labels = []
for i, j in enumerate(os.listdir(input_img_path)):
for evry in os.listdir(os.path.join(input_img_path,j)):
evry_img_path = os.path.join(input_img_path,j,evry)
img_data = tf.gfile.FastGFile(evry_img_path,'rb').read()
img_datas.append(img_data)
img_labels += [i]
return np.array(img_datas),np.array(img_labels)
def load_inception_v3(model_path):
with tf.gfile.FastGFile(model_path,'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
_ = tf.import_graph_def(graph_def=graph_def,name='')
img_datas,img_labels = parse_img(input_img_path)
load_inception_v3(model_path)
batch_size = 500
num_batches = int(len(img_datas) / batch_size)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
tensor = sess.graph.get_tensor_by_name('pool_3/_reshape:0')
for i in range(num_batches):
batch_img_data = img_datas[i*batch_size:(i+1)*batch_size]
batch_img_labels = img_labels[i*batch_size:(i+1)*batch_size]
feature_v = []
for j in batch_img_data:
j_vector = sess.run(tensor,feed_dict={'DecodeJpeg/contents:0':j})
feature_v.append(j_vector)
feature_v = np.vstack(feature_v)
save_path = os.path.join(output_folder,'data_%d.pickle'%i)
with tf.gfile.FastGFile(save_path,'w') as f:
pickle.dump((feature_v,batch_img_labels),f)
print(save_path,'is_ok!')