Update app.py

app.py

@@ -152,32 +152,68 @@
 #         video_capture.release()
 
 
-import torch
-import gradio as gr
+
+
+import streamlit as st
+from ultralytics import YOLO
+import cv2
+import numpy as np
 from PIL import Image
 
-# Load the YOLOv11 model
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', source='local')  # replace 'yolov5' with 'yolo11' if needed
+# Load the YOLO model (replace "best.pt" with your model path if different)
+model = YOLO("best.pt")  
+
+# Streamlit app title
+st.title("Fire Detection in Forest")
+
+# Sidebar for input options
+input_option = st.sidebar.selectbox("Select Input Method", ["Upload Image", "Use Webcam", "Upload Video"])
 
-# Define the prediction function
-def detect_fire(image):
-    # Run the YOLO model on the uploaded image
-    results = model(image)
-    # Check if fire is detected based on the model's output classes
-    detected = any('fire' in results.names[label] for label in results.xyxy[0][:, -1].tolist())
-    return "Fire Detected" if detected else "No Fire Detected"
+if input_option == "Upload Image":
+    # Upload Image
+    uploaded_file = st.file_uploader("Choose an Image", type=["jpg", "jpeg", "png"])
 
-# Create and launch the Gradio interface
-interface = gr.Interface(
-    fn=detect_fire,
-    inputs=gr.Image(type='pil', label="Upload Image"),
-    outputs="text",
-    title="Fire Detection with YOLOv11",
-    description="Upload an image to classify whether fire is detected or not."
-)
+    if uploaded_file is not None:
+        # Open the uploaded image
+        img = Image.open(uploaded_file)
+        st.image(img, caption='Uploaded Image', use_column_width=True)
+        st.write("Classifying...")
 
-interface.launch()
+        # Convert image to a numpy array
+        img_np = np.array(img)
 
+        # Make predictions using the model
+        results = model.predict(source=img_np, conf=0.5)
+
+        # Variable to check if fire is detected
+        fire_detected = False
+
+        # Draw bounding boxes on the image and check for fire detection
+        for result in results:
+            boxes = result.boxes.xyxy.cpu().numpy()
+            class_ids = result.boxes.cls.cpu().numpy().astype(int)
+            
+            for box, class_id in zip(boxes, class_ids):
+                x1, y1, x2, y2 = map(int, box[:4])
+                label = result.names[class_id]
+
+                # Draw bounding box on the image
+                color = (0, 255, 0) if label == "fire" else (255, 0, 0)  # Green for fire, red otherwise
+                img_np = cv2.rectangle(img_np, (x1, y1), (x2, y2), color, 2)
+                img_np = cv2.putText(img_np, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
+
+                # Check if the detected label is "fire"
+                if label.lower() == "fire":
+                    fire_detected = True
+
+        # Show the resulting image with bounding boxes
+        st.image(img_np, caption='Processed Image', use_column_width=True)
+
+        # Display message based on fire detection
+        if fire_detected:
+            st.success("🔥 Fire Detected!")
+        else:
+            st.warning("No Fire Detected.")