keras實現基於孿生網絡的圖片類似度計算

import keras
from keras.layers import Input,Dense,Conv2D
from keras.layers import MaxPooling2D,Flatten,Convolution2D
from keras.models import Model
import os
import numpy as np
from PIL import Image
from keras.optimizers import SGD
from scipy import misc
root_path = os.getcwd()
train_names = ['bear','blackswan','bus','camel','car','cows','dance','dog','hike','hoc','kite','lucia','mallerd','pigs','soapbox','stro','surf','swing','train','walking']
test_names = ['boat','dance-jump','drift-turn','elephant','libby']
 
def load_data(seq_names,data_number,seq_len): 
#生成圖片對
    print('loading data.....')
    frame_num = 51
    train_data1 = []
    train_data2 = []
    train_lab = []
    count = 0
    while count < data_number:
        count = count + 1
        pos_neg = np.random.randint(0,2)
        if pos_neg==0:
           seed1 = np.random.randint(0,seq_len)
           seed2 = np.random.randint(0,seq_len)
           while seed1 == seed2:
             seed1 = np.random.randint(0,seq_len)
             seed2 = np.random.randint(0,seq_len)
           frame1 = np.random.randint(1,frame_num)
           frame2 = np.random.randint(1,frame_num)
           path1 = os.path.join(root_path,'data','simility_data',seq_names[seed1],str(frame1)+'.jpg')
           path2 = os.path.join(root_path, 'data', 'simility_data', seq_names[seed2], str(frame2) + '.jpg')
           image1 = np.array(misc.imresize(Image.open(path1),[224,224]))
           image2 = np.array(misc.imresize(Image.open(path2),[224,224]))
           train_data1.append(image1)
           train_data2.append(image2)
           train_lab.append(np.array(0))
        else:
          seed = np.random.randint(0,seq_len)
          frame1 = np.random.randint(1, frame_num)
          frame2 = np.random.randint(1, frame_num)
          path1 = os.path.join(root_path, 'data', 'simility_data', seq_names[seed], str(frame1) + '.jpg')
          path2 = os.path.join(root_path, 'data', 'simility_data', seq_names[seed], str(frame2) + '.jpg')
          image1 = np.array(misc.imresize(Image.open(path1),[224,224]))
          image2 = np.array(misc.imresize(Image.open(path2),[224,224]))
          train_data1.append(image1)
          train_data2.append(image2)
          train_lab.append(np.array(1))
    return np.array(train_data1),np.array(train_data2),np.array(train_lab)
 
 
def vgg_16_base(input_tensor):
    net = Conv2D(64(3,3),activation='relu',padding='same',input_shape=(224,224,3))(input_tensor)
    net = Convolution2D(64,(3,3),activation='relu',padding='same')(net)
    net = MaxPooling2D((2,2),strides=(2,2))(net)
 
    net = Convolution2D(128,(3,3),activation='relu',padding='same')(net)
    net = Convolution2D(128,(3,3),activation='relu',padding='same')(net)
    net= MaxPooling2D((2,2),strides=(2,2))(net)
 
    net = Convolution2D(256,(3,3),activation='relu',padding='same')(net)
    net = Convolution2D(256,(3,3),activation='relu',padding='same')(net)
    net = Convolution2D(256,(3,3),activation='relu',padding='same')(net)
    net = MaxPooling2D((2,2),strides=(2,2))(net)
 
    net = Convolution2D(512,(3,3),activation='relu',padding='same')(net)
    net = Convolution2D(512,(3,3),activation='relu',padding='same')(net)
    net = Convolution2D(512,(3,3),activation='relu',padding='same')(net)
    net = MaxPooling2D((2,2),strides=(2,2))(net)
 
    net = Convolution2D(512,(3,3),activation='relu',padding='same')(net)
    net = Convolution2D(512,(3,3),activation='relu',padding='same')(net)
    net = Convolution2D(512,(3,3),activation='relu',padding='same')(net)
    net = MaxPooling2D((2,2),strides=(2,2))(net)
    net = Flatten()(net)
    return net
 
 
def siamese(vgg_path=None,siamese_path=None):
    input_tensor = Input(shape=(224,224,3))
    vgg_model = Model(input_tensor,vgg_16_base(input_tensor))
    if vgg_path:
       vgg_model.load_weights(vgg_path)
    input_im1 = Input(shape=(224,224,3))
    input_im2 = Input(shape=(224,224,3))
    out_im1 = vgg_model(input_im1)
    out_im2 = vgg_model(input_im2)
    concatenated = keras.layers.concatenate([out_im1,out_im2])
    out = Dense(500,activation='relu')(concatenated)
    out = Dense(1,activation='sigmoid')(out)
    model = Model([input_im1,input_im2],out)
    if siamese_path:
        model.load_weights(siamese_path)
    return model
 
 
train = True
if train:
   model = siamese(siamese_path='model/simility/vgg.h5')
   sgd = SGD(lr=1e-6,momentum=0.9,decay=1e-6,nesterov=True)
   model.compile(optimizer=sgd,loss='mse',metrics=['accuracy'])
   tensorboard = keras.callbacks.TensorBoard(histogram_freq=5,log_dir='log/simility',write_grads=True,write_images=True)
   ckpt = keras.callbacks.ModelCheckpoint(os.path.join(root_path,'model','simility','vgg.h5'),
                                       verbose=1,period=5)
   train_data1,train_data2,train_lab = load_data(train_names,4000,20)
   model.fit([train_data1,train_data2],train_lab,callbacks=[tensorboard,ckpt],batch_size=64,epochs=50)
else:
   model = siamese(siamese_path='model/simility/vgg.h5')
   test_im1,test_im2,test_labe = load_data(test_names,1000,5)
   TP = 0
   for i in range(1000):
      im1 = np.expand_dims(test_im1[i],axis=0)
      im2 = np.expand_dims(test_im2[i],axis=0)
      lab = test_labe[i]
      pre = model.predict([im1,im2])
      if pre>0.9 and lab==1:
        TP = TP + 1
      if pre<0.9 and lab==0:
        TP = TP + 1
   print(float(TP)/1000)