import os
import tensorflow as tf
from time import time
import VGG16_model as model
import utils
from scipy.misc import imread,imresize
import numpy as np
startTime=time()
batch_size=32
capacity=180
means=[123.68,116.779,103.939]
sess=tf.Session()
init=tf.global_variables_initializer()
sess.run(init)
xs,ys=utils.get_file('./data/train/')
image_batch,label_batch=utils.get_batch(xs,ys,224,224,batch_size,capacity)
x=tf.placeholder(tf.float32,[None,224,224,3])
y=tf.placeholder(tf.int32,[None,2])
vgg=model.vgg16(x)
fc8_finetuining=vgg.probs
loss_function=tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=fc8_finetuining,labels=y))
optimizer=tf.train.GradientDescentOptimizer(0.001).minimize(loss_function)
sess=tf.Session()
init=tf.global_variables_initializer()
sess.run(init)
vgg.load_weights('./data/vgg16_weights.npz',sess)
saver=tf.train.Saver()
print('Model restoting......')
saver.restore(sess,'./model/')
print('traing starts from {} epoch'.format(start+1))
coord=tf.train.Coordinator()
threads=tf.train.start_queue_runners(coord=coord,sess=sess)
epoch_start_time=time()
for i in range(start,1000):
images,labels=sess.run([image_batch,label_batch])
labels=utils.onehot(labels)
sess.run(optimizer,feed_dict={x:images,y:labels})
loss=sess.run(loss_function,feed_dict={x:images,y:labels})
print('现在的损失为:%f'%loss)
epoch_end_time=time()
print('当前训练花费的时间:',(epoch_end_time-epoch_start_time))
epoch_start_time=epoch_end_time
saver.save(sess,ckpt_dir+'cats-vs-dogs_cnn_model.ckpt',global_step=i+1)
sess.run(epoch.assign(i+1))
print('===============Epoch %d is finished==============='%i)
saver.save(sess,'./model/')
print('Optimization Finished!')
coord.request_stop()
coord.join(threads)
filepath='./data/test1/100.jpg'
img=imread(filepath,mode='RGB')
img=imresize(img,(224,224))
img=img.astype(np.float32)
for c in range(3):
img[:,:,c]-=means[c]
prob=sess.run(fc8_finetuining,feed_dict={x:[img]})
max_index=np.argmax(prob)
if max_index==0:
print('这是猫的可能性为:%.6f'%prob[:,0])
else:
print('这是狗的可能性为:%.6f'%prob[:,1])
