# 小型CNN模型搭建和測(cè)試
import picture_read
import numpy as np
import tensorflow as tf
import time

def data_processing(train_document1,train_document2,test_document1,test_document2,    #(訓(xùn)練集數(shù)據(jù)1，訓(xùn)練集數(shù)據(jù)2苫纤，測(cè)試集數(shù)據(jù)1，測(cè)試集數(shù)據(jù)2)
                    pic_size,label1,label2,):    #(圖片縮放大小畴蒲，標(biāo)簽1沮焕，標(biāo)簽2)

    train_data_1,train_label_1 = picture_read.picture_document_deal(train_document1, pic_size, label1)   #sxs標(biāo)簽為0
    train_data_2,train_label_2 = picture_read.picture_document_deal(train_document2, pic_size, label2)   #xsx標(biāo)簽為1
    test_data_1,test_label_1 = picture_read.picture_document_deal(test_document1, pic_size, label1)   #sxs標(biāo)簽為0
    test_data_2,test_label_2 = picture_read.picture_document_deal(test_document2, pic_size, label2)   #xsx標(biāo)簽為1

    # numpy:vstack列上合并兩個(gè)矩陣,hstack在行上合并兩個(gè)矩陣
    train_data = np.vstack((train_data_1,train_data_2))
    train_labels = np.hstack((train_label_1,train_label_2))
    test_data = np.vstack((test_data_1,test_data_2))
    test_labels = np.hstack((test_label_1,test_label_2))
    print('train_data = ',train_data.shape)    #train_data,
    # print('train_labels = ', train_labels, train_labels.shape)

    # 二、 歸一化,將每個(gè)點(diǎn)的像素值歸一化到[0, 1]之間
    training_pics = train_data.reshape(train_data.shape[0], train_data.shape[1], train_data.shape[2], 1)
    test_pics = test_data.reshape(test_data.shape[0], test_data.shape[1], test_data.shape[2], 1)
    training_pics = training_pics / 255.0
    test_pics = test_pics / 255.0
    print('training_pics = ',training_pics.shape)      #training_pics,
    # print('test_pics = ',test_pics,test_pics.shape)

    print(training_pics.shape[1], training_pics.shape[2], training_pics.shape[3])
    print("training_pics = ", training_pics.shape)
    print("training_labels = ", train_labels.shape)

    return training_pics,test_pics,train_labels,test_labels


def model_structure(shape,classification ):   #（圖片的格式隧甚，最終分類的結(jié)果）
    print('shape = ',shape)
    # 第一種寫法
    # 定義模型车荔，按順序創(chuàng)建模型
    model_def = tf.keras.models.Sequential([
        tf.keras.layers.Conv2D(6, (3, 3), activation='relu', input_shape=(shape[1], shape[2], shape[3])),
        tf.keras.layers.MaxPooling2D(2, 2),
        tf.keras.layers.ReLU(),
        tf.keras.layers.Conv2D(16, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D(2, 2),
        tf.keras.layers.ReLU(),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(120, activation='relu'),
        tf.keras.layers.Dense(84, activation='relu'),
        tf.keras.layers.Dense(classification, activation='softmax')
    ])

    # # 第二種寫法
    # inputs = tf.keras.layers.Input(shape = shape[1:])
    # x = tf.keras.layers.Conv2D(filters=6, kernel_size=3, strides=1,
    #                                activation='relu', padding='same')(inputs)
    # x = tf.keras.layers.MaxPool2D(pool_size=2, strides=1, padding='same')(x)
    # x = tf.keras.layers.ReLU()(x)
    # x = tf.keras.layers.Conv2D(filters=16, kernel_size=3, strides=1,
    #                            activation='relu', padding='same')(x)
    # x = tf.keras.layers.MaxPool2D(pool_size=2, strides=1, padding='same')(x)
    # x = tf.keras.layers.ReLU()(x)
    # x = tf.keras.layers.Flatten()(x)
    # x = tf.keras.layers.Dense(120,activation='relu')(x)
    # x = tf.keras.layers.Dense(84,activation='relu')(x)
    # outputs = tf.keras.layers.Dense(classification,activation='softmax')(x)   # 二分類
    # model_def = tf.keras.Model(inputs = inputs, outputs=outputs)


    # 編譯模型，確定優(yōu)化方法戚扳，損失函數(shù)等
    model_def.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

    return model_def

def model_train_save(model,model_file):     #(模型保存文件的名稱)


    # 四忧便、訓(xùn)練模型,并將每輪訓(xùn)練的歷史信息保存在變量 history 中
    history = model.fit(x=training_images, y=train_label, batch_size = 32 ,
                        epochs = 20,validation_split=0.2)  # validation_split = 0.1   validation_data=(test_images, test_labels)
    # 計(jì)算損失
    test_loss = model.evaluate(test_images, test_label)
    # print('test_loss = ',test_loss)

    # 模型保存
    model.save(model_file)    # 保存全部

    # 刪除模型
    del model

def model_predict_labels(model_file,test_pic,test_labels):  #測(cè)試集有標(biāo)簽時(shí)使用
    # 四、恢復(fù)模型并預(yù)測(cè)
    # print("模型讀回")
    model = tf.keras.models.load_model(model_file)
    # 測(cè)試模型
    loss, acc = model.evaluate(test_pic, test_labels)

    # 五帽借、預(yù)測(cè)
    predict = model.predict(test_pic)
    # print('predict = ',predict)
    predict = np.array(tf.argmax(predict, 1))
    print('predict = ',predict)   #tf.argmax根據(jù)axis取值的不同返回每行或者每列最大值的索引,axis = 1時(shí)取每行最大值的索引位置
    print('test_labels = ', test_labels)

    print("Restored model, accuracy:{:5.2f}%".format(100 * acc))

    return predict

def model_predict_nolabels(model_file,test_pic):    #測(cè)試集無標(biāo)簽時(shí)使用
    # 四珠增、恢復(fù)模型
    # print("模型讀回")
    model = tf.keras.models.load_model(model_file)

    # 五、預(yù)測(cè)
    predict = model.predict(test_pic)
    # print('predict = ',predict)
    predict = np.array(tf.argmax(predict, 1))
    print('predict = ',predict)   #tf.argmax根據(jù)axis取值的不同返回每行或者每列最大值的索引,axis = 1時(shí)取每行最大值的索引位置

    return predict

if __name__ == '__main__':

    start_time = time.time()

    # 一砍艾、獲取數(shù)據(jù)并處理數(shù)據(jù)
    im_size = (28, 28)   #圖片放縮的大小
    train_document_1 = 'D:\\python_work\\model_project\\dataset\\train\\sxs'
    train_document_2 = 'D:\\python_work\\model_project\\dataset\\train\\xsx'
    test_document_1 = 'D:\\python_work\\model_project\\dataset\\test\\sxs'
    test_document_2 = 'D:\\python_work\\model_project\\dataset\\test\\xsx'

    training_images, test_images, train_label, test_label = data_processing(train_document_1,train_document_2,
                                                                            test_document_1,test_document_2,   #(訓(xùn)練集數(shù)據(jù)1蒂教，訓(xùn)練集數(shù)據(jù)2，測(cè)試集數(shù)據(jù)1辐董，測(cè)試集數(shù)據(jù)2)
                                                                            im_size,0,1,)
    # 二悴品、模型搭建:
    model_name = model_structure(training_images.shape,classification = 2)

    # 三禀综、訓(xùn)練神經(jīng)網(wǎng)絡(luò)模型進(jìn)行圖片分類:
    model_train_save(model_name,'my_model.h5')


    # 四简烘、模型讀回及預(yù)測(cè)（有對(duì)比標(biāo)簽）
    predict_labels = model_predict_labels('my_model.h5',test_images,test_label)

    # # 五、模型讀回及預(yù)測(cè)（無對(duì)比標(biāo)簽）
    # predict_labels = model_predict_nolabels('my_model.h5',test_images)


    end_time = time.time()
    print('程序運(yùn)行時(shí)間：',end_time-start_time)