import os
import tensorflow as tf
from PIL import Image
import numpy as np


from memryx import NeuralCompiler, Simulator

# Get model
URL = "http://download.tensorflow.org/models/mobilenet_v1_2018_02_22"
TARGET = "mobilenet_v1_1.0_224"
os.system("wget  {}/{}.tgz".format(URL, TARGET))
os.system("tar -xvf {}.tgz".format(TARGET))

# Prepare image
image = Image.open('image.png').resize((224,224))
image = np.array(image)
image = image / np.max(image) * 2 - 1
image = np.expand_dims(image, axis=0)

# Compile model
dfp = NeuralCompiler(models=TARGET+"_frozen.pb", verbose=1).run()

# Run Simulator
s = Simulator(dfp=dfp, verbose=1)
outputs = s.infer(inputs=image)
latency, fps = s.benchmark(frames=4)

# Process outputs and save the classification result!
outputs = np.expand_dims(outputs[0], 0)
outputs = outputs[:, 1:]
predictions = tf.keras.applications.mobilenet.decode_predictions(outputs, top=1)[0][0]
print("I see a '{}' with {:.1f} % certainty".format(predictions[1], predictions[2]*100))
print("Simulated MXA FPS: ", fps)
print("Simulated MXA Latency: ", latency)