import gradio as gr
from ultralytics import YOLO
import numpy as np
import cv2

model = YOLO("best.pt")


def detect(image):
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    results = model.predict(img, conf=0.25, verbose=False)

    annotated = results[0].plot()

    annotated = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)

    text = []

    for box in results[0].boxes:
        cls = int(box.cls.item())
        conf = float(box.conf.item())

        text.append(f"{model.names[cls]} : {conf:.2%}")

    if len(text) == 0:
        text = ["✅ SAFE — No Fire or Smoke Detected"]

    return annotated, "\n".join(text)


demo = gr.Interface(
    fn=detect,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(label="Detection"),
        gr.Textbox(label="Prediction")
    ],
    title="Fire & Smoke Detection"
)

demo.launch()