#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#Original code can be found at
#https://colab.research.google.com/github/tensorflow/examples/blob/master/community/en/flowers_tf_lite.ipynb#scrollTo=aCLb_yV5JfF3

import tensorflow as tf
import os
import numpy as np
import matplotlib.pyplot as plt

IMAGE_SIZE = 224
BATCH_SIZE = 64

def download_flower_dataset():
    _URL = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
    zip_file = tf.keras.utils.get_file(origin=_URL,
                                        fname="flower_photos.tgz",
                                        extract=True)
    return os.path.join(os.path.dirname(zip_file), 'flower_photos')


def create_image_batch_generator(base_dir):
    datagen = tf.keras.preprocessing.image.ImageDataGenerator(
        rescale=1./255,
        validation_split=0.2)

    train_generator = datagen.flow_from_directory(
        base_dir,
        target_size=(IMAGE_SIZE, IMAGE_SIZE),
        batch_size=BATCH_SIZE,
        subset='training')

    val_generator = datagen.flow_from_directory(
        base_dir,
        target_size=(IMAGE_SIZE, IMAGE_SIZE),
        batch_size=BATCH_SIZE,
        subset='validation')
    
    return train_generator, val_generator

 
def save_labels(train_generator):
    for image_batch, label_batch in train_generator:
        break

    print(image_batch.shape, label_batch.shape)
    print (train_generator.class_indices)

    labels = '\n'.join(sorted(train_generator.class_indices.keys()))

    with open('labels.txt', 'w') as f:
        f.write(labels)


def download_mobilenet_v2_model():
    # Create the base model from the pre-trained model MobileNet V2
    IMG_SHAPE = (IMAGE_SIZE, IMAGE_SIZE, 3)
    base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
                                                include_top=False,
                                                weights='imagenet')

    model = tf.keras.Sequential([
        base_model,
        tf.keras.layers.Conv2D(32, 3, activation='relu'),
        tf.keras.layers.Dropout(0.2), 
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dense(5, activation='softmax')
    ])

    # Let's take a look to see how many layers are in the base model
    print("Number of layers in the base model: ", len(base_model.layers))

    return base_model, model

def run_transfer_learning(base_model, model, train_generator, val_generator):
    base_model.trainable = False
    model.compile(optimizer=tf.keras.optimizers.Adam(),
                loss='categorical_crossentropy',
                metrics=['accuracy'])

    model.summary()
    print('Number of trainable variables = {}'.format(len(model.trainable_variables)))

    epochs = 10
    history = model.fit(train_generator,
                        epochs=epochs,
                        validation_data=val_generator)
    return history

def run_fine_tuning(base_model, model, train_generator, val_generator):
    base_model.trainable = True
    # Fine tune from this layer onwards
    fine_tune_at = 100

    # Freeze all the layers before the `fine_tune_at` layer
    for layer in base_model.layers[:fine_tune_at]:
        layer.trainable =  False

    model.compile(loss='categorical_crossentropy',
                optimizer = tf.keras.optimizers.Adam(1e-5),
                metrics=['accuracy'])
    model.summary()

    print('Number of trainable variables = {}'.format(len(model.trainable_variables)))

    history = model.fit(train_generator,
                            epochs=5,
                            validation_data=val_generator)
    return history

def save_model_as_tflite(model):
    saved_model_dir = 'fine_tuning'

    tf.saved_model.save(model, saved_model_dir)
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
    tflite_model = converter.convert()

    with open('model.tflite', 'wb') as f:
        f.write(tflite_model)

def plot_figure(history, fig_name):
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']

    plt.figure(figsize=(8, 8))
    plt.subplot(2, 1, 1)
    plt.plot(acc, label='Training Accuracy')
    plt.plot(val_acc, label='Validation Accuracy')
    plt.legend(loc='lower right')
    plt.ylabel('Accuracy')
    plt.ylim([min(plt.ylim()),1])
    plt.title('Training and Validation Accuracy')

    plt.subplot(2, 1, 2)
    plt.plot(loss, label='Training Loss')
    plt.plot(val_loss, label='Validation Loss')
    plt.legend(loc='upper right')
    plt.ylabel('Cross Entropy')
    plt.ylim([0,1.0])
    plt.title('Training and Validation Loss')
    plt.xlabel('epoch')
    plt.show()
    plt.savefig(fig_name)
    
if __name__ == '__main__':
    print(tf.__version__)
    base_dir = download_flower_dataset()
    train_generator, val_generator = create_image_batch_generator(base_dir)
    save_labels(train_generator)

    base_model, model = download_mobilenet_v2_model() #download without top layer and add top layer 

    history = run_transfer_learning(base_model, model, train_generator, val_generator)
    plot_figure(history, 'transfer_learning.png')

    history_fine = run_fine_tuning(base_model, model, train_generator, val_generator)
    save_model_as_tflite(model)
    plot_figure(history_fine, 'fine_tuning.png')


matplotlib.pyplot.legend matplotlib.pyplot.plot tensorflow.keras.preprocessing.image.ImageDataGenerator tensorflow.keras.layers.Conv2D tensorflow.saved_model.save matplotlib.pyplot.subplot tensorflow.lite.TFLiteConverter.from_saved_model matplotlib.pyplot.figure matplotlib.pyplot.title matplotlib.pyplot.ylim tensorflow.keras.layers.Dense matplotlib.pyplot.savefig matplotlib.pyplot.xlabel matplotlib.pyplot.show matplotlib.pyplot.ylabel tensorflow.keras.applications.MobileNetV2 tensorflow.keras.layers.GlobalAveragePooling2D tensorflow.keras.optimizers.Adam tensorflow.keras.utils.get_file tensorflow.keras.layers.Dropout
[(26, 'tensorflow.keras.utils.get_file', 'tf.keras.utils.get_file', ([], {'origin': '_URL', 'fname': '"""flower_photos.tgz"""', 'extract': '(True)'}), True, 'import tensorflow as tf\n'), (33, 'tensorflow.keras.preprocessing.image.ImageDataGenerator', 'tf.keras.preprocessing.image.ImageDataGenerator', ([], {'rescale': '(1.0 / 255)', 'validation_split': '(0.2)'}), True, 'import tensorflow as tf\n'), (68, 'tensorflow.keras.applications.MobileNetV2', 'tf.keras.applications.MobileNetV2', ([], {'input_shape': 'IMG_SHAPE', 'include_top': '(False)', 'weights': '"""imagenet"""'}), True, 'import tensorflow as tf\n'), (124, 'tensorflow.saved_model.save', 'tf.saved_model.save', (['model', 'saved_model_dir'], {}), True, 'import tensorflow as tf\n'), (125, 'tensorflow.lite.TFLiteConverter.from_saved_model', 'tf.lite.TFLiteConverter.from_saved_model', (['saved_model_dir'], {}), True, 'import tensorflow as tf\n'), (137, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(8, 8)'}), True, 'import matplotlib.pyplot as plt\n'), (138, 'matplotlib.pyplot.subplot', 'plt.subplot', (['(2)', '(1)', '(1)'], {}), True, 'import matplotlib.pyplot as plt\n'), (139, 'matplotlib.pyplot.plot', 'plt.plot', (['acc'], {'label': '"""Training Accuracy"""'}), True, 'import matplotlib.pyplot as plt\n'), (140, 'matplotlib.pyplot.plot', 'plt.plot', (['val_acc'], {'label': '"""Validation Accuracy"""'}), True, 'import matplotlib.pyplot as plt\n'), (141, 'matplotlib.pyplot.legend', 'plt.legend', ([], {'loc': '"""lower right"""'}), True, 'import matplotlib.pyplot as plt\n'), (142, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Accuracy"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (144, 'matplotlib.pyplot.title', 'plt.title', (['"""Training and Validation Accuracy"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (146, 'matplotlib.pyplot.subplot', 'plt.subplot', (['(2)', '(1)', '(2)'], {}), True, 'import matplotlib.pyplot as plt\n'), (147, 'matplotlib.pyplot.plot', 'plt.plot', (['loss'], {'label': '"""Training Loss"""'}), True, 'import matplotlib.pyplot as plt\n'), (148, 'matplotlib.pyplot.plot', 'plt.plot', (['val_loss'], {'label': '"""Validation Loss"""'}), True, 'import matplotlib.pyplot as plt\n'), (149, 'matplotlib.pyplot.legend', 'plt.legend', ([], {'loc': '"""upper right"""'}), True, 'import matplotlib.pyplot as plt\n'), (150, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Cross Entropy"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (151, 'matplotlib.pyplot.ylim', 'plt.ylim', (['[0, 1.0]'], {}), True, 'import matplotlib.pyplot as plt\n'), (152, 'matplotlib.pyplot.title', 'plt.title', (['"""Training and Validation Loss"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (153, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""epoch"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (154, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (155, 'matplotlib.pyplot.savefig', 'plt.savefig', (['fig_name'], {}), True, 'import matplotlib.pyplot as plt\n'), (29, 'os.path.dirname', 'os.path.dirname', (['zip_file'], {}), False, 'import os\n'), (74, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['(32)', '(3)'], {'activation': '"""relu"""'}), True, 'import tensorflow as tf\n'), (75, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['(0.2)'], {}), True, 'import tensorflow as tf\n'), (76, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'tf.keras.layers.GlobalAveragePooling2D', ([], {}), True, 'import tensorflow as tf\n'), (77, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(5)'], {'activation': '"""softmax"""'}), True, 'import tensorflow as tf\n'), (87, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {}), True, 'import tensorflow as tf\n'), (110, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', (['(1e-05)'], {}), True, 'import tensorflow as tf\n'), (143, 'matplotlib.pyplot.ylim', 'plt.ylim', ([], {}), True, 'import matplotlib.pyplot as plt\n')]