# --- directory structure ---
# image_classification.py
# data:
#   class1:
#       1.jpeg
#       2.jpeg
#       ...
#   class2:
#       1.jpeg
#       2.jpeg
#       ...
#   class3:
#       1.jpeg
#       xyz.jpeg
#       ...

from keras.applications.inception_v3 import InceptionV3
from keras.preprocessing import image
from keras.models import Model
from keras.layers import GlobalAveragePooling2D, Dense, Dropout, Activation, Flatten
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras.utils import np_utils
from keras import backend as K

import cv2
import skimage
import matplotlib.pyplot as plt

import numpy as np

import os, time
import glob

os.environ['KMP_DUPLICATE_LIB_OK']='True'

def load_normalized_image(filename, image_size):
    return (skimage.transform.resize(plt.imread(filename).astype(float), [image_size, image_size]) - 255.0/2) / 255.0

def generator(data_root, images_list, labels_list, batch_size, image_size, class_names):
    num_classes = len(class_names)
    batch_features = np.zeros((batch_size, image_size, image_size, 3))

    while True:
        labels = []

        for i in range(batch_size):
            index = np.random.choice(len(images_list))
            image_path = os.path.join(data_root, images_list[index])
            img = load_normalized_image(image_path, image_size)
            if len(img.shape) < 3:
                img = cv2.merge((img, img, img))
            if len(img.shape) > 3:
                img = img[:, :, 0:3]
            batch_features[i] = img
            labels.append(labels_list[index])

        batch_labels = np_utils.to_categorical(labels, num_classes)
        yield batch_features, batch_labels

def load_data(data_dirs, filetypes):
    data = []
    labels = []
    for idx, data_dir in enumerate(data_dirs):
        files = []
        for filetype in filetypes:
            files.extend(glob.glob(os.path.join(data_dir, '*.' + filetype)))
        data.append(files)
        labels.append(np.ones(len(data[idx]), dtype=int) * idx)

    return data, labels

data_directory = 'data'
filetypes = ['jpeg', 'jpg', 'png']

checkpoints_directory = 'checkpoints'
log_directory = 'logs'

class_names = [os.path.basename(d) for d in glob.glob(os.path.join(data_directory, '*')) if os.path.isdir(d)]
data_dirs = [os.path.join(data_directory, class_name) for class_name in class_names]

# Split data into training and valudation sets

data, labels = load_data(data_dirs, filetypes)

x_train = []
y_train = []
x_test = []
y_test = []
for i in range(len(data)):
    x_train.extend(data[i][:160])
    y_train.extend(labels[i][:160])

    x_test.extend(data[i][160:])
    y_test.extend(labels[i][160:])

# Hyper parameters
image_size = 224
image_channels = 3
batch_size = 8
epochs = 10

input_shape = (image_size, image_size, image_channels)
num_classes = len(data)

model = InceptionV3(weights='imagenet', include_top=False, input_shape=input_shape)
# model.summary()

for layer in model.layers:
    layer.trainable = False

# Adding custom Layers
x = model.output
x = Flatten()(x)
x = Dense(4096, activation="relu")(x)
x = Dropout(0.5)(x)
x = Dense(4096, activation="relu")(x)
x = Dropout(0.5)(x)
predictions = Dense(num_classes, activation="softmax")(x)

# Creating the final model
model_final = Model(input=model.input, output=predictions)
# model_final.summary()

# Compile model
model_final.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy'])

# Callbacks
if not os.path.exists(checkpoints_directory):
    os.makedirs(checkpoints_directory)

chkpoint_callback = ModelCheckpoint(
    os.path.join(checkpoints_directory, 'weights.{epoch:02d}.hdf5'),
    monitor = 'val_acc',
    verbose = 1,
    save_best_only = False,
    save_weights_only = False,
    mode = 'auto',
    period = 10)

tb_callback = TensorBoard(
    log_dir = log_directory,
    histogram_freq = 0,
    batch_size = batch_size,
    write_graph = True,
    write_images = True)

hist = model_final.fit_generator(
    generator('./', x_train, y_train, batch_size, image_size, class_names),
    steps_per_epoch = len(x_train),
    epochs = epochs,
    verbose = 1,
    validation_data = generator('./', x_test, y_test, batch_size, image_size, class_names),
    validation_steps = len(x_test),
    callbacks = [chkpoint_callback, tb_callback])