Update app.py

app.py

@@ -5,13 +5,12 @@ import numpy as np
 from PIL import Image
 
 # Load the YOLO model (replace with your model path)
-model = YOLO("best.pt")
-   # Use your YOLO model file here
+model = YOLO("best.pt")  # Use your YOLO model file here
 
 st.title("Fire Detection in Forest")
 
 # Sidebar for input options
-input_option = st.sidebar.selectbox("Select Input Method", ["Upload Image", "Use Webcam","Upload Video"])
+input_option = st.sidebar.selectbox("Select Input Method", ["Upload Image", "Use Webcam", "Upload Video"])
 
 if input_option == "Upload Image":
     # Upload Image
@@ -27,21 +26,32 @@ if input_option == "Upload Image":
 
         # Make predictions
         results = model.predict(source=img_np, conf=0.5)
-        
-        if Fire_Detected= True:
-            print("Fire Detected")
-        else:
-            print("Not")
+
+        # Variable to check if fire is detected
+        fire_detected = False
+
         # Draw bounding boxes on the image
         for result in results:
             boxes = result.boxes.xyxy
             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
         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.")
+
+
 # elif input_option == "Use Webcam":
 #     st.write("Starting webcam for live detection...")