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

model = YOLO('hotspot_detector.pt')

def detect_hotspots(image):
    result = model(image)
    
    for r in result:
        im_array = r.plot()
        # im = Image.fromarray(im_array[..., ::-1])
        return Image.fromarray(im_array[..., ::-1])


css = """
.gr-example img {
    width: 300px !important;
    height: auto !important;
}
"""

demo = gr.Interface(fn=detect_hotspots, inputs=gr.Image(type='pil'), outputs="image", title="YOLO Object Detection",
                    examples=[['images/test_image_1.jpg'], ['images/test_image_2.jpg'],
                              ['images/test_image_3.jpg'], ['images/test_image_4.jpg']],
                    examples_per_page=4,
                    cache_examples= False,
                    css=css
                    )  
demo.launch()