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


def preprocess_image_bytes(image):
    image = Image.open(io.BytesIO(image))
    image = np.array(image)
    image = np.expand_dims(image, axis=0)

    return image


# Import Image and Model as bytes (Simulate call from API)
benign_image_bytes = open("data/test/benign/5.jpg", "rb").read()
malignant_image_bytes = open("data/test/malignant/1.jpg", "rb").read()
model_bytes = open("models/mobilenet_v3.keras", "rb")


# Convert model to Keras model
model_file = open("model.keras", "wb").write(model_bytes.read())
model = tf.keras.models.load_model("model.keras")
os.remove("model.keras")


benign_image = preprocess_image_bytes(benign_image_bytes)
malignant_image = preprocess_image_bytes(malignant_image_bytes)

benign_prediction = model.predict(benign_image)
malignant_prediction = model.predict(malignant_image)


print(f"Benign Prediction: {benign_prediction}")
print(f"Malignant Prediction: {malignant_prediction}")