from typing import IO

import tensorflow as tf
from gradio import Interface, inputs, outputs

SIZE = 224

inter = tf.lite.Interpreter("lite_model.tflite", num_threads=12)
inter.allocate_tensors()
in_tensor, *_ = inter.get_input_details()
out_tensor, *_ = inter.get_output_details()


def process_data(content):
    img = tf.io.decode_jpeg(content, channels=3)
    img = tf.image.resize_with_pad(img, SIZE, SIZE, method="nearest")
    img = tf.image.resize(img, (SIZE, SIZE), method="nearest")
    img = img / 255
    return img


def main(file: IO[bytes]):
    data = process_data(file.read())
    data = tf.expand_dims(data, 0)
    inter.set_tensor(in_tensor["index"], data)
    inter.invoke()
    result, *_ = inter.get_tensor(out_tensor["index"])
    safe, questionable, explicit = map(float, result)
    possibilities = {"safe": safe, "questionable": questionable, "explicit": explicit}
    print("Predict result:", possibilities)
    return possibilities


image = inputs.Image(type="file")
label = outputs.Label(type="confidences")

interface = Interface(main, image, label)
interface.launch()