import math

import gradio as gr
import tensorflow as tf

configs = [
    {
        "model": "my_model_2.h5", "size": 512
    },
    {
        "model": "my_model.h5", "size": 224
    },
]

config = configs[0]

new_model = tf.keras.models.load_model(config["model"])

def classify_image(inp):
    inp = inp.reshape((-1, config["size"], config["size"], 3))
    prediction = new_model.predict(inp).flatten()
    print(prediction)
    if len(prediction) > 1:
        probability = 100 * math.exp(prediction[0]) / (math.exp(prediction[0]) + math.exp(prediction[1]))
    else:
        probability = round(100. / (1 + math.exp(-prediction[0])), 2)
    if probability > 45:
        return "Glaucoma", probability
    if probability > 25:
        return "Unclear", probability
    return "Not glaucoma", probability


gr.Interface(
    fn=classify_image, 
    inputs=gr.inputs.Image(shape=(config["size"], config["size"])),
    outputs=[
    gr.outputs.Textbox(label="Label"),
    gr.outputs.Textbox(label="Glaucoma probability (0 - 100)"),
    ],
    examples=["001.jpg", "002.jpg", "225.jpg"],
    flagging_options=["Correct label", "Incorrect label"],
    allow_flagging="manual",
).launch()