import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO

# Load the YOLO model
model = YOLO("best.pt")  # Ensure the path to your model is correct

# Set the title of the app
st.title("Live Fire Detection App")

# Create a placeholder for the video stream
video_placeholder = st.empty()

# Function to perform live fire detection
def detect_fire():
    # Open a connection to the webcam (0 is the default camera)
    cap = cv2.VideoCapture(0)

    while True:
        ret, frame = cap.read()
        if not ret:
            st.error("Failed to capture video.")
            break

        # Convert the frame to RGB (as YOLO expects RGB input)
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Make predictions on the current frame
        results = model.predict(source=rgb_frame, conf=0.5)

        # Draw bounding boxes on the frame
        fire_detected = False
        for result in results:
            boxes = result.boxes.xyxy  # Bounding boxes
            for box in boxes:
                x1, y1, x2, y2 = box[:4].astype(int)
                rgb_frame = cv2.rectangle(rgb_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)

                # Check if the detected class is "fire" (adjust based on your model's class mapping)
                class_id = int(box[5])  # Assuming class ID is at the 6th position
                if class_id == 0:  # Replace 0 with the actual class ID for fire if different
                    fire_detected = True

        # Display the frame in the Streamlit app
        video_placeholder.image(rgb_frame, channels="RGB", use_column_width=True)

        # Display message based on fire detection
        if fire_detected:
            st.success("🔥 Fire Detected!")
        else:
            st.warning("No Fire Detected.")

        # Break the loop if the user presses 'q'
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # Release the webcam and close windows
    cap.release()
    cv2.destroyAllWindows()

# Start the detection process
if st.button("Start Live Detection"):
    detect_fire()