import gradio as gr
from ultralytics import YOLO
from PIL import Image

CONF_THRESH=0.332
IOU=0.4

model = YOLO("best.pt")

def predict(image):
    results = model(image, verbose=False, conf=CONF_THRESH, iou=IOU, agnostic_nms=True)
    results[0].names = {0: 'bee', 1: 'stylopidae', 2: 'meloidae larvae'}
    bgr_result_img = results[0].plot()
    rgb_result_img = Image.fromarray(bgr_result_img[..., ::-1])
    return rgb_result_img

def main():
    with gr.Blocks() as demo:
        gr.HTML("<h1>Deep Learning Based Detection of Bee Parasites under Natural Conditions</h1>")

        with gr.Row():
            image_input = gr.Image(type="pil", label="Upload Image")
            image_output = gr.Image(label="Prediction")

        gr.Examples(
            examples=["stylops.jpeg", "larvae.jpeg"],
            inputs=image_input
        )

        submit_btn = gr.Button("Detect")
        submit_btn.click(fn=predict, inputs=image_input, outputs=[image_output])

        gr.HTML("""
                  <div>
                    Sources:
                    <ul>
                      <li><a href="https://www.inaturalist.org/observations/265792162">Stylops</a></li>
                      <li><a href="https://www.inaturalist.org/observations/278157533">Meloidae larvae</a></li>
                    </ul>
                  </div>
                  <img src="https://upload.wikimedia.org/wikipedia/commons/thumb/9/98/TU_Ilmenau_Logo_black_green.svg/208px-TU_Ilmenau_Logo_black_green.svg.png" alt="TU Ilmenau" style="position: absolute; bottom: 0; right: 0; width: 20vw; max-width: 250px; height: auto;">
                """)

    demo.launch()

if __name__ == '__main__':
    main()