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



description = """
<center><img src="https://huggingface.co./spaces/intelliarts/hotspot-anomaly-detection-for-solar-panels/resolve/main/images/ia_logo.png" width=270px> </center><br>
<p style="font-size: 20px; text-align: center;"">This is a demo of a computer vision model by Intelliarts. It's designed to detect anomalies in solar panels, in particular, overheated spots. The model operates on infrared images of solar panels. It indicates the overheated spot, outlines its area, and shows the approximate accuracy of anomaly detection. You can use your own infrared images for testing or utilize samples from our dataset. This demo is not a finished ML solution, but rather a proof of concept. It can be extended to detect other anomalies in solar panels as well as visible damages.</p>
"""


demo = gr.Interface(fn=detect_hotspots, inputs=gr.Image(type='pil'), outputs="image",
                    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,
                    description=description
                    )  
demo.launch()