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(), outputs="image", title="YOLO Object Detection")
demo.launch()