from ultralytics import YOLO
import streamlit as st
import numpy as np
from PIL import Image
model=YOLO("best.pt")
# Set the title of the app
st.title("Fire Detection App")

# Sidebar for input options
input_option = st.sidebar.selectbox("Select Input Method", ["Upload Image"])

if input_option == "Upload Image":
    # Upload Image
    uploaded_file = st.file_uploader("Choose an Image", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        # Open and display the uploaded image
        img = Image.open(uploaded_file)
        st.image(img, caption='User Image', use_column_width=True)
        
        # Convert the image to a numpy array
        img_np = np.array(img)

        # Make predictions
        results = model.predict(source=img_np, conf=0.5)  # Adjust confidence threshold as needed

        # Variable to check if fire is detected
        fire_detected = False

        # Draw bounding boxes on the image
        for result in results:
            boxes = result.boxes.xyxy  # Bounding boxes
            for box in boxes:
                x1, y1, x2, y2 = box[:4].astype(int)
                img_np = cv2.rectangle(img_np, (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

        # Show the resulting image with bounding boxes
        st.image(img_np, caption='Detected Fire', use_column_width=True)

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