Update app.py

app.py

@@ -9,40 +9,61 @@ model = YOLO("best.pt")  # Ensure the path to your model is correct
 # Set the title of the app
 st.title("Live Fire Detection App")
 
-def capture_video():
-    """Attempts to capture video from the default camera (index 0).
+# Load the video capture
+videoCap = cv2.VideoCapture(0)
 
-    Returns:
-        cv2.VideoCapture: The video capture object if successful, None otherwise.
-    """
+# Function to get class colors
+def getColours(cls_num):
+    base_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]
+    color_index = cls_num % len(base_colors)
+    increments = [(1, -2, 1), (-2, 1, -1), (1, -1, 2)]
+    color = [base_colors[color_index][i] + increments[color_index][i] * 
+    (cls_num // len(base_colors)) % 256 for i in range(3)]
+    return tuple(color)
 
-    cap = cv2.VideoCapture(0)
 
-    if not cap.isOpened():
-        print("Error: Could not open video capture device.")
-        return None
+while True:
+    ret, frame = videoCap.read()
+    if not ret:
+        continue
+    results = yolo.track(frame, stream=True)
 
-    return cap
 
-# Capture video
-cap = capture_video()
-if cap is None:
-    # Handle capture failure gracefully
-    exit()
+    for result in results:
+        # get the classes names
+        classes_names = result.names
 
-while True:
-    ret, frame = cap.read()
+        # iterate over each box
+        for box in result.boxes:
+            # check if confidence is greater than 40 percent
+            if box.conf[0] > 0.4:
+                # get coordinates
+                [x1, y1, x2, y2] = box.xyxy[0]
+                # convert to int
+                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
 
-    if not ret:
-        print("Error: Could not read frame.")
-        break
+                # get the class
+                cls = int(box.cls[0])
+
+                # get the class name
+                class_name = classes_names[cls]
+
+                # get the respective colour
+                colour = getColours(cls)
+
+                # draw the rectangle
+                cv2.rectangle(frame, (x1, y1), (x2, y2), colour, 2)
 
-    # Process the frame here (e.g., display, save, analyze)
-    cv2.imshow('Frame', frame)
+                # put the class name and confidence on the image
+                cv2.putText(frame, f'{classes_names[int(box.cls[0])]} {box.conf[0]:.2f}', (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1, colour, 2)
+                
+    # show the image
+    cv2.imshow('frame', frame)
 
+    # break the loop if 'q' is pressed
     if cv2.waitKey(1) & 0xFF == ord('q'):
         break
 
-# Release resources
-cap.release()
-cv2.destroyAllWindows()
+# release the video capture and destroy all windows
+videoCap.release()
+cv2.destroyAllWindows()