import numpy as np
from keras.datasets import mnist
import gc

from keras.models import Sequential, Model
from keras.layers import Input, Dense, Dropout, Flatten
from keras.layers.convolutional import Conv2D, MaxPooling2D
from keras.applications.vgg16 import VGG16
from keras.optimizers import SGD
from keras.utils.data_utils import get_file
import cv2
import h5py as h5py 
import numpy as np

import os
import math
from matplotlib import pyplot as plt

#全局变量
RATIO = 0.2
train_dir = 'D:/dl4cv/datesets/littleCBIR/'

#根据分类总数确定one-hot总类
NUM_DENSE = 5
#训练总数
epochs = 10

def tran_y(y): 
    y_ohe = np.zeros(NUM_DENSE) 
    y_ohe[y] = 1 
    return y_ohe

#根据Ratio获得训练和测试数据集的图片地址和标签
##生成数据集,本例先验3**汽车、4**恐龙、5**大象、6**花、7**马
def get_files(file_dir, ratio):
    '''
    Args:
        file_dir: file directory
    Returns:
        list of images and labels
    '''
    image_list = []
    label_list = []
    for file in os.listdir(file_dir):
        if file[0:1]=='3':
            image_list.append(file_dir + file)
            label_list.append(0)
        elif file[0:1]=='4':
            image_list.append(file_dir + file)
            label_list.append(1)
        elif file[0:1]=='5':
            image_list.append(file_dir + file)
            label_list.append(2)
        elif file[0:1]=='6':
            image_list.append(file_dir + file)
            label_list.append(3)
        else:
            image_list.append(file_dir + file)
            label_list.append(4)
    print('数据集导入完毕')
    #图片list和标签list
    #hstack 水平(按列顺序)把数组给堆叠起来
    image_list = np.hstack(image_list)
    label_list = np.hstack(label_list)
    
    temp = np.array([image_list, label_list])
    temp = temp.transpose()
    np.random.shuffle(temp)   
    
    all_image_list = temp[:, 0]
    all_label_list = temp[:, 1]
    
    n_sample = len(all_label_list)
    #根据比率，确定训练和测试数量
    n_val = math.ceil(n_sample*ratio) # number of validation samples
    n_train = n_sample - n_val # number of trainning samples
    tra_images = []
    val_images = []
    #按照0-n_train为tra_images，后面位val_images的方式来排序
    
    for index in range(n_train):
        image = cv2.imread(all_image_list[index])
        #灰度，然后缩放
        image = cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
        image = cv2.resize(image,(48,48))#到底在这个地方修改，还是在后面修改，需要做具体实验
        tra_images.append(image)

    tra_labels = all_label_list[:n_train]
    tra_labels = [int(float(i)) for i in tra_labels]

    for index in range(n_val):
        image = cv2.imread(all_image_list[n_train+index])
        #灰度，然后缩放
        image = cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
        image = cv2.resize(image,(32,32))
        val_images.append(image)

    val_labels = all_label_list[n_train:]
    val_labels = [int(float(i)) for i in val_labels]
    return np.array(tra_images),np.array(tra_labels),np.array(val_images),np.array(val_labels)

# colab+VGG要求至少48像素在现有数据集上，已经能够完成不错情况
ishape=48
#(X_train, y_train), (X_test, y_test) = mnist.load_data() 
#获得数据集
#X_train, y_train, X_test, y_test = get_files(train_dir, RATIO)
#保持数据
##np.savez("D:\\dl4cv\\datesets\\littleCBIR.npz",X_train=X_train,y_train=y_train,X_test=X_test,y_test=y_test)
#读取数据
 
path='littleCBIR.npz'
#https://github.com/jsxyhelu/GOCW/raw/master/littleCBIR.npz
path = get_file(path,origin='https://github.com/jsxyhelu/GOCW/raw/master/littleCBIR.npz')
f = np.load(path)
X_train, y_train = f['X_train'], f['y_train']
X_test, y_test = f['X_test'], f['y_test']

X_train = [cv2.cvtColor(cv2.resize(i, (ishape, ishape)), cv2.COLOR_GRAY2BGR) for i in X_train] 
X_train = np.concatenate([arr[np.newaxis] for arr in X_train]).astype('float32') 
X_train /= 255.0

X_test = [cv2.cvtColor(cv2.resize(i, (ishape, ishape)), cv2.COLOR_GRAY2BGR) for i in X_test] 
X_test = np.concatenate([arr[np.newaxis] for arr in X_test]).astype('float32')
X_test /= 255.0

y_train_ohe = np.array([tran_y(y_train[i]) for i in range(len(y_train))]) 
y_test_ohe = np.array([tran_y(y_test[i]) for i in range(len(y_test))])
y_train_ohe = y_train_ohe.astype('float32')
y_test_ohe = y_test_ohe.astype('float32')

model_vgg = VGG16(include_top = False, weights = 'imagenet', input_shape = (ishape, ishape, 3)) 
#for i, layer in enumerate(model_vgg.layers): 
#    if i<20:
for layer in model_vgg.layers:
        layer.trainable = False
model = Flatten()(model_vgg.output) 
model = Dense(4096, activation='relu', name='fc1')(model)
model = Dense(4096, activation='relu', name='fc2')(model)
model = Dropout(0.5)(model)
model = Dense(NUM_DENSE, activation = 'softmax', name='prediction')(model) 
model_vgg_pretrain = Model(model_vgg.input, model, name = 'vgg16_pretrain')
#model_vgg_pretrain.summary()
print("vgg准备完毕\n")
sgd = SGD(lr = 0.05, decay = 1e-5) 
model_vgg_pretrain.compile(loss = 'categorical_crossentropy', optimizer = sgd, metrics = ['accuracy'])
print("vgg开始训练\n")
log = model_vgg_pretrain.fit(X_train, y_train_ohe, validation_data = (X_test, y_test_ohe), epochs = epochs, batch_size = 64)

score = model_vgg_pretrain.evaluate(X_test, y_test_ohe, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])

plt.figure('acc')  
plt.subplot(2, 1, 1)  
plt.plot(log.history['acc'],'r--',label='Training Accuracy')  
plt.plot(log.history['val_acc'],'r-',label='Validation Accuracy')  
plt.legend(loc='best')  
plt.xlabel('Epochs')  
plt.axis([0, epochs, 0.5, 1])  
plt.figure('loss')  
plt.subplot(2, 1, 2)  
plt.plot(log.history['loss'],'b--',label='Training Loss')  
plt.plot(log.history['val_loss'],'b-',label='Validation Loss')  
plt.legend(loc='best')  
plt.xlabel('Epochs')  
plt.axis([0, epochs, 0, 1])  
  
plt.show() 
os.system("pause")