Update app.py

app.py

@@ -156,64 +156,37 @@
 
 import streamlit as st
 from ultralytics import YOLO
-import cv2
 import numpy as np
 from PIL import Image
 
-# Load the YOLO model (replace "best.pt" with your model path if different)
-model = YOLO("best.pt")  
+# Load the YOLO model (use the path to your 'best.pt' file)
+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"])
+# Upload Image
+uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
 
-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 the uploaded image
+    img = Image.open(uploaded_file)
+    st.image(img, caption="Uploaded Image", use_column_width=True)
 
-    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...")
+    # Convert image to a numpy array
+    img_np = np.array(img)
 
-        # Convert image to a numpy array
-        img_np = np.array(img)
+    # Make predictions
+    results = model.predict(source=img_np, imgsz=640, conf=0.5)
 
-        # Make predictions using the model
-        results = model.predict(source=img_np, conf=0.5)
+    # Check if fire is detected
+    fire_detected = any("fire" in results.names[int(cls)] for cls in results[0].boxes.cls)
 
-        # Variable to check if fire is detected
-        fire_detected = False
+    # Display results
+    if fire_detected:
+        st.success("🔥 Fire Detected!")
+    else:
+        st.warning("No Fire Detected.")
 
-        # 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.")