Update app.py

app.py

@@ -1,34 +1,44 @@
 import cv2
 import numpy as np
+import torch
 import streamlit as st
+from ultralytics import YOLO
 from camera_input_live import camera_input_live
 
-# Load Haarcascade for face detection
-cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")
+# Load YOLO fire detection model
+model_path = "last.pt"
+if not torch.cuda.is_available():
+    device = "cpu"
+else:
+    device = "cuda"
+
+model = YOLO(model_path)
+model.to(device)
 
 # Streamlit app title
-st.title("Live Object Detection with Camera")
-st.subheader("Hold your face in front of the webcam to see real-time detection.")
+st.title("Live Fire Detection with Camera")
+st.subheader("Hold the camera towards potential fire sources to detect in real-time.")
 
 # Capture live camera input
 image = camera_input_live()
 
 if image is not None:
-
     # Convert the image to OpenCV format
     bytes_data = image.getvalue()
     cv2_img = cv2.imdecode(np.frombuffer(bytes_data, np.uint8), cv2.IMREAD_COLOR)
 
-    # Convert to grayscale for face detection
-    gray = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2GRAY)
-
-    # Detect faces in the image
-    faces = cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=3)
+    # Perform fire detection
+    results = model(cv2_img, device=device)
 
-    # Draw rectangles around detected faces
-    for (x, y, w, h) in faces:
-        cv2.rectangle(cv2_img, (x, y), (x + w, y + h), (0, 255, 0), 3)
-        cv2.putText(cv2_img, "Face", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+    # Draw bounding boxes for detected fires
+    for result in results:
+        boxes = result.boxes
+        for box in boxes:
+            b = box.xyxy[0].cpu().numpy().astype(int)
+            c = int(box.cls[0])
+            label = f'Fire {box.conf[0]:.2f}'
+            cv2.rectangle(cv2_img, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 3)
+            cv2.putText(cv2_img, label, (b[0], b[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
 
     # Display the annotated image
-    st.image(cv2_img, channels="BGR", caption="Detected Faces", use_container_width=True)
+    st.image(cv2_img, channels="BGR", caption="Detected Fire", use_container_width=True)