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]#VGG训练时图像预处理所减均值(RGB三通道)
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#即sofemax(fc8)微调（finetuining）sofemax(fc8)
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)#GradientDescentOptimizer下降优化器
sess=tf.Session()#声明会话
init=tf.global_variables_initializer()#调用变量
sess.run(init)#运行变量
vgg.load_weights('./data/vgg16_weights.npz',sess)#通过npz格式的文件获取VGG的相应权重参数，从而将权重注入即可实现复用
saver=tf.train.Saver()#生成saver，用于保存和提取变量
print('Model restoting......')
saver.restore(sess,'./model/')#恢复最后保存的模型
#saver.restore(sess,'.model/epoch_00800.ckpt')恢复指定检查点的模型
print('traing starts from {} epoch'.format(start+1))
coord=tf.train.Coordinator()#使用协调器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)#用one-hot对标签进行编码
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)#join操作等待其他线程结束，其他所有线程关闭之后，这一函数才能返回
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])
