Update app.py

app.py

@@ -9,61 +9,19 @@ model = YOLO("best.pt")  # Ensure the path to your model is correct
 # Set the title of the app
 st.title("Live Fire Detection App")
 
-# Load the video capture
-videoCap = cv2.VideoCapture(0)
-
-# 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)
-
+webcamera = cv2.VideoCapture(0)
+# webcamera.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
+# webcamera.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
 
 while True:
-    ret, frame = videoCap.read()
-    if not ret:
-        continue
-    results = yolo.track(frame, stream=True)
-
-
-    for result in results:
-        # get the classes names
-        classes_names = result.names
-
-        # 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)
-
-                # 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)
-
-                # 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)
+    success, frame = webcamera.read()
+    
+    results = model.track(frame, classes=0, conf=0.8, imgsz=480)
+    cv2.putText(frame, f"Total: {len(results[0].boxes)}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
+    cv2.imshow("Live Camera", results[0].plot())
 
-    # break the loop if 'q' is pressed
-    if cv2.waitKey(1) & 0xFF == ord('q'):
+    if cv2.waitKey(1) == ord('q'):
         break
 
-# release the video capture and destroy all windows
-videoCap.release()
+webcamera.release()
 cv2.destroyAllWindows()