import gradio as gr
from model import predict
from io import BytesIO


def predict_image(image, model_type):
    """
    image: PIL Image
    model_type: 'variety' or 'disease'
    """
    buf = BytesIO()
    image.save(buf, format="PNG")
    image_bytes = buf.getvalue()

    return predict(image_bytes, model_type=model_type)


gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(type="pil", label="Upload Banana Image"),
        gr.Radio(
            choices=["variety", "disease"],
            value="variety",
            label="Select Model"
        )
    ],
    outputs=gr.JSON(label="Prediction Result"),
).launch()