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])


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']])
demo.launch()