app.py

import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
import pandas as pd
import pickle
import random
import torch.serialization

# Load the Random Forest model directly from the Space repository
try:
    with open("model.pkl", "rb") as f:
        rf_model = pickle.load(f)
except Exception as e:
    raise Exception(f"Failed to load Random Forest model: {str(e)}")

# Load the YOLO model with safe globals directly from the Space repository
try:
    from ultralytics.nn.tasks import DetectionModel
    import torch.nn as nn
    torch.serialization.add_safe_globals([DetectionModel, nn.Sequential])
    yolo_model = YOLO("best.pt")
except Exception as e:
    raise Exception(f"Failed to load YOLO model: {str(e)}")

# Simulated sensor data (from simulated_sensors.py)
def simulate_sensor_data(scenario="Safe"):
    if scenario == "Safe":
        data = {
            "temperature": random.uniform(20, 30),
            "humidity": random.uniform(30, 50),
            "mq2_smoke": random.uniform(0, 100),
            "mq135_gas": random.uniform(0, 100),
            "flame_detected": False
        }
    elif scenario == "Gas Leak":
        data = {
            "temperature": random.uniform(20, 30),
            "humidity": random.uniform(30, 50),
            "mq2_smoke": random.uniform(100, 300),
            "mq135_gas": random.uniform(600, 1000),
            "flame_detected": False
        }
    elif scenario == "Fire Detected":
        data = {
            "temperature": random.uniform(40, 50),
            "humidity": random.uniform(30, 50),
            "mq2_smoke": random.uniform(300, 600),
            "mq135_gas": random.uniform(100, 300),
            "flame_detected": True
        }
    elif scenario == "Warning":
        data = {
            "temperature": random.uniform(30, 40),
            "humidity": random.uniform(50, 70),
            "mq2_smoke": random.uniform(200, 400),
            "mq135_gas": random.uniform(300, 500),
            "flame_detected": False
        }
    elif scenario == "Evacuate Immediately":
        data = {
            "temperature": random.uniform(45, 50),
            "humidity": random.uniform(30, 50),
            "mq2_smoke": random.uniform(600, 1000),
            "mq135_gas": random.uniform(400, 600),
            "flame_detected": True
        }
    else:
        data = {
            "temperature": 0,
            "humidity": 0,
            "mq2_smoke": 0,
            "mq135_gas": 0,
            "flame_detected": False
        }
    return data

# YOLO inference function
def detect_fire_smoke_from_image(image):
    if image is None:
        return {"cv_flame_score": 0.0, "cv_smoke_score": 0.0, "person_detected": 0}, image

    results = yolo_model(image)
    flame_score = 0.0
    smoke_score = 0.0
    person_detected = 0

    for result in results:
        for box in result.boxes:
            label = result.names[int(box.cls)]
            conf = float(box.conf)

            if "fire" in label.lower():
                flame_score = max(flame_score, conf)
            elif "smoke" in label.lower():
                smoke_score = max(smoke_score, conf)
            elif "person" in label.lower():
                person_detected = 1

            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, f"{label} {conf:.2f}", (x1, y1 - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return {
        "cv_flame_score": round(flame_score, 3),
        "cv_smoke_score": round(smoke_score, 3),
        "person_detected": person_detected
    }, image

# Combined prediction function
def predict_system(scenario, image=None):
    # Simulate sensor data
    sensor_data = simulate_sensor_data(scenario)
    
    # Perform YOLO inference if an image is provided
    if image is not None:
        vision_data, annotated_image = detect_fire_smoke_from_image(image)
    else:
        vision_data = {"cv_flame_score": 0.0, "cv_smoke_score": 0.0, "person_detected": 0}
        annotated_image = np.zeros((480, 640, 3), dtype=np.uint8)
        cv2.putText(annotated_image, "No Image Uploaded", (50, 240), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

    # Combine sensor and vision data
    combined_data = {**sensor_data, **vision_data}

    # Prepare features for Random Forest model
    features = pd.DataFrame([[
        combined_data["temperature"],
        combined_data["humidity"],
        combined_data["mq2_smoke"],
        combined_data["mq135_gas"],
        1 if combined_data["flame_detected"] else 0,
        combined_data["cv_flame_score"],
        combined_data["cv_smoke_score"],
        1 if combined_data["person_detected"] else 0
    ]], columns=[
        "temperature", "humidity", "mq2_smoke", "mq135_gas",
        "flame_detected", "cv_flame_score", "cv_smoke_score", "person_detected"
    ])

    # Predict threat level
    prediction = rf_model.predict(features)[0]

    # Format the output
    output = f"""
    **Threat Level:** {prediction}

    **Sensor Data:**
    - Temperature: {combined_data["temperature"]:.2f} °C
    - Humidity: {combined_data["humidity"]:.2f} %
    - MQ2 Smoke: {combined_data["mq2_smoke"]:.2f} ppm
    - MQ135 Gas: {combined_data["mq135_gas"]:.2f} ppm
    - Flame Detected: {"Yes" if combined_data["flame_detected"] else "No"}

    **Vision Data:**
    - CV Flame Score: {combined_data["cv_flame_score"] * 100:.2f}%
    - CV Smoke Score: {combined_data["cv_smoke_score"] * 100:.2f}%
    - Person Detected: {"Yes" if combined_data["person_detected"] else "No"}
    """
    return output, annotated_image

# Create Gradio interface
inputs = [
    gr.Dropdown(choices=["Safe", "Gas Leak", "Fire Detected", "Warning", "Evacuate Immediately"], label="Scenario", value="Safe"),
    gr.Image(type="numpy", label="Upload Image for Vision Detection (Optional)")
]

outputs = [
    gr.Textbox(label="System Output"),
    gr.Image(type="numpy", label="Annotated Image")
]

gr.Interface(
    fn=predict_system,
    inputs=inputs,
    outputs=outputs,
    title="Fire & Gas Leak Detection System",
    description="Simulate sensor data and upload an image for vision detection to predict the threat level using a Random Forest model and YOLOv10."
).launch()