import tensorflow as tf
import os, shutil

MODEL_LOCATION='export/flowers_model3'  # will be created
# load from checkpoint and export a model that has desired signature
CHECK_POINT_DIR='gs://practical-ml-vision-book/flowers_5_trained/chkpts'
model = tf.keras.models.load_model(CHECK_POINT_DIR)

IMG_HEIGHT = 345
IMG_WIDTH = 345
IMG_CHANNELS = 3
CLASS_NAMES = 'daisy dandelion roses sunflowers tulips'.split()
    
# a single image of any size
@tf.function(input_signature=[tf.TensorSpec([None, None, 3], dtype=tf.float32)])
def predict_flower_type(img):
    img = tf.image.resize_with_pad(img, IMG_HEIGHT, IMG_WIDTH)
    batch_pred = model(tf.expand_dims(img, axis=0))
    top_prob = tf.math.reduce_max(batch_pred, axis=[1])
    pred_label_index = tf.math.argmax(batch_pred, axis=1)
    pred_label = tf.gather(tf.convert_to_tensor(CLASS_NAMES), pred_label_index)
    return {
        'probability': tf.squeeze(top_prob, axis=0),
        'flower_type': tf.squeeze(pred_label, axis=0)
    }

shutil.rmtree('export', ignore_errors=True)
os.mkdir('export')


model.save(MODEL_LOCATION,
          signatures={
              'serving_default': predict_flower_type
          })

"""
Explanation: Edge ML with TensorFlow Lite
In this notebook, we convert the saved model into a TensorFlow Lite model
so that we can run it on Edge devices.
In order to do edge inference, we need to handle raw image data from the camera
and process a single image (not a batch of images).
End of explanation
"""


import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_saved_model(MODEL_LOCATION)
tflite_model = converter.convert()

with open('export/model.tflite', 'wb') as ofp:
    ofp.write(tflite_model)

!ls -lh export/model.tflite

"""
Explanation: Convert to TFLite
This will take a while to do the conversion
End of explanation
"""


