"""
Simple example on how to fine tune the VGG16 model in Keras and how to use them.
Part 1

Source: https://www.guru99.com/keras-tutorial.html
"""

#### Data preparation
from keras.preprocessing.image import ImageDataGenerator
import numpy as np
import matplotlib.pyplot as plt
import keras
from keras.layers import Flatten, Dense
from keras.applications.vgg16 import VGG16
from keras.optimizers import SGD

TRAIN_PATH  = 'images/train/'
TEST_PATH   = 'images/test/'
BATCH_SIZE  = 16
IMAGE_SIZE  = 224
NUM_CLASSES = 8

# The ImageDataGenerator will make an X_training data from a directory.
# The sub-directory in that directory will be used as a class for each object.
# The image will be loaded with the RGB color mode, with the categorical class
# mode for the Y_training data, with a batch size of 16. Finally, shuffle the
# data.

train_datagen = ImageDataGenerator(validation_split=0.3,
                                   shear_range=0.2,
                                   zoom_range=0.2,
                                   horizontal_flip=True)

train_generator = train_datagen.flow_from_directory(
                        directory=TRAIN_PATH,
                        target_size=(IMAGE_SIZE,IMAGE_SIZE),
                        batch_size=BATCH_SIZE,
                        class_mode='categorical',
                        color_mode='rgb',
                        shuffle=True)


# Let's see our images randomly by plotting them with matplotlib
# x_batch, y_batch = train_generator.next()
x_batch, y_batch = next(train_generator)
fig=plt.figure()

COLUMNS = 4
ROWS    = 4

for i in range(1, COLUMNS*ROWS):
    num = np.random.randint(BATCH_SIZE)
    image = x_batch[num].astype(np.int)
    fig.add_subplot(ROWS, COLUMNS, i)
    plt.imshow(image)
plt.show()

### Creating model
# Let's create our network model from VGG16 with imageNet pre-trained weight.
# We will freeze these layers so that the layers are not trainable to help us
# reduce the computation time.

# Load the VGG model
base_model = VGG16(weights='imagenet', include_top=False, input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3))

print(base_model.summary())

# Freeze the layers
for layer in base_model.layers:
    layer.trainable = False

# Create the model
model = keras.models.Sequential()

# Add the vgg convolutional base model
model.add(base_model)

# Add new layers
model.add(Flatten())
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(NUM_CLASSES, activation='softmax'))

# Show a summary of the model. Check the number of trainable parameters
print(model.summary())

input("Press Enter to continue...")

# As you can see, the summary of our network model. From an input from
# VGG16 Layers, then we add 2 Fully Connected Layer which will extract 1024
# features and an output layer that will compute the 8 classes with the softmax
# activation.

#### Training
# Compile the model
model.compile(loss='categorical_crossentropy',
              optimizer=SGD(lr=1e-3),
              metrics=['accuracy'])

# Start the training process
# model.fit(x_train, y_train, validation_split=0.30, batch_size=32, epochs=50, verbose=2)

# Save the model
# model.save('catdog.h5')

history = model.fit_generator(train_generator,
                              steps_per_epoch=train_generator.n/BATCH_SIZE,
                              epochs=10)

model.save('fine_tune.h5')

# Summarize history for accuracy
plt.plot(history.history['loss'])
plt.title('loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['loss'], loc='upper left')
plt.show()

# Our losses are dropped significantly and the accuracy is almost 100%. For
# testing our model, we randomly picked images over the internet and put it on
# the test folder with a different class to test