from tensorflow import keras
from PIL import Image
import numpy as np

from memryx import NeuralCompiler, Simulator

# Get model
model = keras.applications.MobileNet()

# Prepare image
image = Image.open('image.png').resize((224,224))
image = keras.applications.mobilenet.preprocess_input(np.array(image))
image = np.expand_dims(image, axis=0)

# Compile model
dfp = NeuralCompiler(models=model, 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)
predictions = 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)