Update app.py

app.py

@@ -3,11 +3,52 @@ from ultralytics import YOLO
 import cv2
 from deep_sort_realtime.deepsort_tracker import DeepSort
 
+def initialize_tracker(max_age=30, n_init=3, nn_budget=100):
+    return DeepSort(max_age=max_age, n_init=n_init, nn_budget=nn_budget)
+
+def detect_people(model, frame, confidence_threshold=0.5):
+    results = model(frame, device="cpu")  # Force CPU if CUDA is unavailable
+    detections = []
+    for result in results:
+        for box, cls, conf in zip(result.boxes.xyxy, result.boxes.cls, result.boxes.conf):
+            if result.names[int(cls)] == "person" and conf > confidence_threshold:
+                x1, y1, x2, y2 = map(int, box)
+                bbox = [x1, y1, x2 - x1, y2 - y1]
+                detections.append((bbox, conf, "person"))
+    return detections
+
+def count_people_in_video(video_path, model_path="setosys_ppl_in_video_small_v1.pt", confidence_threshold=0.5):
+    logging.basicConfig(level=logging.INFO)
+    cap = cv2.VideoCapture(video_path)
+    model = YOLO(model_path)  # Auto-detect device during model loading
+    tracker = initialize_tracker()
+
+    total_ids = set()
+    frame_count = 0
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        frame_count += 1
+        logging.info(f"Processing frame {frame_count}")
+        detections = detect_people(model, frame, confidence_threshold)
+        tracks = tracker.update_tracks(detections, frame=frame)
+
+        for track in tracks:
+            if track.is_confirmed():
+                total_ids.add(track.track_id)
+
+    cap.release()
+    return len(total_ids)
+
+
 # Initialize YOLO model
 model = YOLO("setosys_ppl_in_video_small_v1.pt")  # Load  model
-tracker = DeepSort(max_age=30, n_init=3, nn_budget=100)
+#tracker = DeepSort(max_age=30, n_init=3, nn_budget=100)
 
-def count_people_in_video(video_path):
+def count_people_in_video_old(video_path):
     cap = cv2.VideoCapture(video_path)  # Load video
     total_ids = set()  # Track unique IDs