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 ] )