from tensorflow.keras.applications.vgg19 import VGG19, preprocess_input
from tensorflow.keras.models import Model
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense
from tensorflow.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
from tensorflow.keras.optimizers import SGD
import tensorflow.keras.backend as K
# 訓練和測試的圖片分爲'bus', 'dinosaur', 'flower', 'horse', 'elephant'五類
# 其圖片的下載地址爲 http://pan.baidu.com/s/1nuqlTnN ,總共500張圖片,其中圖片以3,4,5,6,7開頭進行按類區分
# 訓練圖片400張，測試圖片100張；注意下載後，在train和test目錄下分別建立上述的五類子目錄，keras會按照子目錄進行分類識別
NUM_CLASSES = 5
TRAIN_PATH = '/home/yourname/Documents/tensorflow/images/500pics/train'
TEST_PATH = '/home/yourname/Documents/tensorflow/images/500pics/test'
# 代碼最後挑出一張圖片進行預測識別
PREDICT_IMG = '/home/yourname/Documents/tensorflow/images/500pics/test/elephant/502.jpg'
# FC層定義輸入層的大小
FC_NUMS = 1024
# 凍結訓練的層數，根據模型的不同，層數也不一樣，根據調試的結果，VGG19和VGG16c層比較符合理想的測試結果，本文采用VGG19做示例
FREEZE_LAYERS = 17
# 進行訓練和測試的圖片大小，VGG19推薦爲224×244
IMAGE_SIZE = 224
# 採用VGG19爲基本模型，include_top爲False，表示FC層是可自定義的，拋棄模型中的FC層；該模型會在~/.keras/models下載基本模型
base_model = VGG19(input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3), include_top=False, weights='imagenet')
# 自定義FC層以基本模型的輸入爲卷積層的最後一層
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(FC_NUMS, activation='relu')(x)
prediction = Dense(NUM_CLASSES, activation='softmax')(x)
# 構造完新的FC層，加入custom層
model = Model(inputs=base_model.input, outputs=prediction)
# 可觀察模型結構
model.summary()
# 獲取模型的層數
print("layer nums:", len(model.layers))
# 除了FC層，靠近FC層的一部分卷積層可參與參數訓練，
# 一般來說，模型結構已經標明一個卷積塊包含的層數，
# 在這裏我們選擇FREEZE_LAYERS爲17，表示最後一個卷積塊和FC層要參與參數訓練
for layer in model.layers[:FREEZE_LAYERS]:
layer.trainable = False
for layer in model.layers[FREEZE_LAYERS:]:
layer.trainable = True
for layer in model.layers:
print("layer.trainable:", layer.trainable)
# 預編譯模型
model.compile(optimizer=SGD(lr=0.0001, momentum=0.9), loss='categorical_crossentropy', metrics=['accuracy'])
# 給出訓練圖片的生成器， 其中classes定義後，可讓model按照這個順序進行識別
train_datagen = ImageDataGenerator()
train_generator = train_datagen.flow_from_directory(directory=TRAIN_PATH,
target_size=(IMAGE_SIZE, IMAGE_SIZE), classes=['bus', 'dinosaur', 'flower', 'horse', 'elephant'])
test_datagen = ImageDataGenerator()
test_generator = test_datagen.flow_from_directory(directory=TEST_PATH,
target_size=(IMAGE_SIZE, IMAGE_SIZE), classes=['bus', 'dinosaur', 'flower', 'horse', 'elephant'])
# 運行模型
model.fit_generator(train_generator, epochs=5, validation_data=test_generator)
# 找一張圖片進行預測驗證
img = load_img(path=PREDICT_IMG, target_size=(IMAGE_SIZE, IMAGE_SIZE))
# 轉換成numpy數組
x = img_to_array(img)
# 轉換後的數組爲3維數組(224,224,3),
# 而訓練的數組爲4維(圖片數量, 224,224, 3),所以我們可擴充下維度
x = K.expand_dims(x, axis=0)
# 需要被預處理下
x = preprocess_input(x)
# 數據預測
result = model.predict(x, steps=1)
# 最後的結果是一個含有5個數的一維數組，我們取最大值所在的索引號，即對應'bus', 'dinosaur', 'flower', 'horse', 'elephant'的順序
print("result:", K.eval(K.argmax(result)))