import gradio as gr
import yolov5
from PIL import Image

app_title = "Aksara Jawa Layout Detection"
models_id = 'hermanshid/yolo-layout-detector'

model = yolov5.load(models_id)

examples = [['test_images/example1.jpg', 0.6, ]]


def predict(image, threshold=0.25):
    global models_id
    global model

    input_size = 640

    model.conf = threshold
    results = model(image, size=input_size)
    numpy_image = results.render()[0]
    output_image = Image.fromarray(numpy_image)
    return output_image


gr.Interface(
    title=app_title,
    description="Created by 'hermanshid'",
    fn=predict,
    inputs=[
        gr.Image(type="pil"),
        gr.Slider(maximum=1, step=0.01, value=0.25),
    ],
    outputs=gr.Image(type="pil"),
    examples=examples,
    cache_examples=True if examples else False,
).launch(enable_queue=True)