import gradio as gr
import cv2
from ultralytics import YOLO

# 1. Load your new masterpiece model
model = YOLO("best.pt")

def predict_image(img):
    if img is None:
        return None, "No image uploaded."
    
    # Convert Gradio's RGB format to BGR for YOLO
    bgr_img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    # Run prediction (Balanced threshold at 0.25 confidence)
    results = model.predict(source=bgr_img, conf=0.25)
    
    # Get the visually annotated BGR image and map back to RGB for Gradio
    annotated_img_bgr = results[0].plot()
    annotated_img_rgb = cv2.cvtColor(annotated_img_bgr, cv2.COLOR_BGR2RGB)
    
    # Extract detected classes safely using native, pre-mapped indices
    detected_classes = []
    if results[0].boxes is not None:
        for box in results[0].boxes:
            cls_id = int(box.cls[0])
            class_name = model.names[cls_id] # Natively tracks 'Smoke' or 'Fire' perfectly
            detected_classes.append(class_name)
    
    # Generate the warning message
    if len(detected_classes) == 0:
        status_warning = "✅ SYSTEM STATUS: Safe (No Fire or Smoke detected)"
    else:
        unique_threats = list(set(detected_classes))
        threats_str = " & ".join(unique_threats)
        status_warning = f"🚨 WARNING: {threats_str.upper()} DETECTED!"

    return annotated_img_rgb, status_warning

# Build the Masterpiece Gradio UI Layout
with gr.Blocks(title="🔥 AI Fire & Smoke Detection System v2.0") as demo:
    gr.Markdown("# 🔥 AI Fire & Smoke Detection System v2.0 (Masterpiece Edition)")
    gr.Markdown("An advanced custom-trained YOLOv8 system optimized against glare, ambient lighting, and complex vapor patterns.")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="numpy", label="Upload CCTV Frame")
            submit_btn = gr.Button("Analyze Frame", variant="primary")
        with gr.Column():
            output_img = gr.Image(type="numpy", label="Detections")
            status_output = gr.Textbox(label="System Warning / Alert Status", interactive=False)
            
    submit_btn.click(
        fn=predict_image,
        inputs=input_img,
        outputs=[output_img, status_output]
    )

if __name__ == "__main__":
    demo.launch()