"""
Simple example on how to fine tune models in Keras and how to use them.
Part 2

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

from keras.models import load_model
from keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import numpy as np

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

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

### Testing our model
model = load_model('fine_tune.h5')

test_datagen = ImageDataGenerator()
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)

test_generator = test_datagen.flow_from_directory(
                        directory=TEST_PATH,
                        target_size=(IMAGE_SIZE, IMAGE_SIZE),
                        color_mode='rgb',
                        shuffle=False,
                        class_mode='categorical',
                        batch_size=1)

filenames  = test_generator.filenames
nb_samples = len(filenames)

fig = plt.figure()

COLUMNS  = 4
ROWS     = 4

# And our test is as given below! Only 1 image is predicted wrong from a test of
# 14 images!

for i in range(1, COLUMNS*ROWS -1):
    x_batch, y_batch = next(test_generator)

    name = model.predict(x_batch)
    name = np.argmax(name, axis=-1)
    true_name = y_batch
    true_name = np.argmax(true_name, axis=-1)

    label_map = (test_generator.class_indices)
    label_map = dict((v,k) for k,v in label_map.items()) #flip k,v
    predictions = [label_map[k] for k in name]
    true_value = [label_map[k] for k in true_name]

    image = x_batch[0].astype(np.int)
    fig.add_subplot(ROWS, COLUMNS, i)
    plt.title(str(predictions[0]) + ':' + str(true_value[0]))
    plt.imshow(image)
plt.show()