## from https://medium.com/@14prakash/transfer-learning-using-keras-d804b2e04ef8

print("[INFO] Loading necessary libraries.")
from keras import applications
from keras.preprocessing.image import ImageDataGenerator
from keras import optimizers
from keras.models import Sequential, Model 
from keras.layers import Dropout, Flatten, Dense, GlobalAveragePooling2D
from keras import backend as k 
from keras.callbacks import ModelCheckpoint, LearningRateScheduler, TensorBoard, EarlyStopping
from time import time

print("[INFO] Establishing the location and size of our images.")
img_width, img_height = 256, 256
train_data_dir = "/home/ubuntu/Research/HST_Public_DLN/Data/train"
validation_data_dir = "/home/ubuntu/Research/HST_Public_DLN/Data/validation"
nb_train_samples = 4125
nb_validation_samples = 466 

print("[INFO] Establishing the run parameters for the network.")
batch_size = 16
epochs = 50

print("[INFO] Creating image augmentation generators.")
# Initiate the train and test generators with data Augumentation 
print('[INFO]', end=" ")
train_datagen = ImageDataGenerator(
rescale = 1./255,
horizontal_flip = True,
fill_mode = "nearest",
zoom_range = 0.3,
width_shift_range = 0.3,
height_shift_range=0.3,
rotation_range=30)

print('[INFO]', end=" ")
test_datagen = ImageDataGenerator(
rescale = 1./255,
horizontal_flip = True,
fill_mode = "nearest",
zoom_range = 0.3,
width_shift_range = 0.3,
height_shift_range=0.3,
rotation_range=30)

print("[INFO] Creating flow control for image access.")
train_generator = train_datagen.flow_from_directory(
train_data_dir,
target_size = (img_height, img_width),
batch_size = batch_size, 
class_mode = "categorical")

validation_generator = test_datagen.flow_from_directory(
validation_data_dir,
target_size = (img_height, img_width),
class_mode = "categorical")

print("[INFO] Creatin gour set of call back operations.")
# Save the model according to the conditions  
tensboard = TensorBoard(log_dir='./logs/log-{}'.format(int(time())))
checkpoint = ModelCheckpoint("vgg16_1.h5", monitor='val_acc', verbose=1, save_best_only=True, save_weights_only=False, mode='auto', period=1)
early = EarlyStopping(monitor='val_acc', min_delta=0, patience=10, verbose=1, mode='auto')

print("[INFO] Establishing the base model network to transfer from.")
model = applications.VGG19(weights = "imagenet", include_top=False, input_shape = (img_width, img_height, 3))

print("[INFO] Turning off all layers except top layer.")
# Freeze the layers which you don't want to train. Here I am freezing the first 5 layers.
for layer in model.layers:#[:5]
    layer.trainable = False

#Adding custom Layers 
print("[INFO] Adding new layers for transfer flexibility.")

hidden_new_layer1 = 1024
hidden_new_layer2 = 1024
dropout_rate = 0.5

x = model.output
x = Flatten()(x)
x = Dense(hidden_new_layer1, activation="relu")(x)
x = Dropout(dropout_rate)(x)
x = Dense(hidden_new_layer2, activation="relu")(x)
predictions = Dense(train_generator.num_classes, activation="softmax")(x)

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

print("[INFO] Compiling the model for transfer training.")
# compile the model 
model_final.compile(loss = "categorical_crossentropy", optimizer='adam', metrics=["accuracy"])

print(model_final.summary())

print("[INFO] Fitting the transfer network with `fit_generator` flowing from directory.")
# Train the model 
model_final.fit_generator(
train_generator,
samples_per_epoch = nb_train_samples,
epochs = epochs,
validation_data = validation_generator,
nb_val_samples = nb_validation_samples,
callbacks = [checkpoint, early, tensboard])

print("[INFO] Finished full run process")

"""
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         (None, 256, 256, 3)       0         
_________________________________________________________________
block1_conv1 (Conv2D)        (None, 256, 256, 64)      1792      
_________________________________________________________________
block1_conv2 (Conv2D)        (None, 256, 256, 64)      36928     
_________________________________________________________________
block1_pool (MaxPooling2D)   (None, 128, 128, 64)      0         
_________________________________________________________________
block2_conv1 (Conv2D)        (None, 128, 128, 128)     73856     
_________________________________________________________________
block2_conv2 (Conv2D)        (None, 128, 128, 128)     147584    
_________________________________________________________________
block2_pool (MaxPooling2D)   (None, 64, 64, 128)       0         
_________________________________________________________________
block3_conv1 (Conv2D)        (None, 64, 64, 256)       295168    
_________________________________________________________________
block3_conv2 (Conv2D)        (None, 64, 64, 256)       590080    
_________________________________________________________________
block3_conv3 (Conv2D)        (None, 64, 64, 256)       590080    
_________________________________________________________________
block3_conv4 (Conv2D)        (None, 64, 64, 256)       590080    
_________________________________________________________________
block3_pool (MaxPooling2D)   (None, 32, 32, 256)       0         
_________________________________________________________________
block4_conv1 (Conv2D)        (None, 32, 32, 512)       1180160   
_________________________________________________________________
block4_conv2 (Conv2D)        (None, 32, 32, 512)       2359808   
_________________________________________________________________
block4_conv3 (Conv2D)        (None, 32, 32, 512)       2359808   
_________________________________________________________________
block4_conv4 (Conv2D)        (None, 32, 32, 512)       2359808   
_________________________________________________________________
block4_pool (MaxPooling2D)   (None, 16, 16, 512)       0         
_________________________________________________________________
block5_conv1 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_conv2 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_conv3 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_conv4 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_pool (MaxPooling2D)   (None, 8, 8, 512)         0         
=================================================================
Total params: 20,024,384.0
Trainable params: 20,024,384.0
Non-trainable params: 0.0
"""
