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

# Load YOLOv10 Fire + Smoke model
model = YOLO("best.pt") 

def detect_fire_smoke(image):
    if image is None:
        return "Please upload an image"

    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    results = model(img)[0]

    if len(results.boxes) == 0:
        return "✔ SAFE — No Fire or Smoke Detected"

    output = []

    for box in results.boxes:
        cls_id = int(box.cls[0])  # YOLOv10 classes
        conf = float(box.conf[0])

        if cls_id == 0:
            output.append(f"🔥 FIRE DETECTED — Confidence {conf:.2f}")
        elif cls_id == 1:
            output.append(f"💨 SMOKE DETECTED — Confidence {conf:.2f}")

    if not output:
        return "✔ SAFE — No Fire or Smoke Detected"

    return "\n".join(output)

demo = gr.Interface(
    fn=detect_fire_smoke,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Fire & Smoke Detection (YOLOv10)",
    description="Upload an image to detect fire or smoke."
)

demo.launch()