Update detect_people.py

detect_people.py

@@ -1,13 +1,104 @@
+import os
 import cv2
+import numpy as np
+from pathlib import Path
+import urllib.request
 
-def detect_people(image_path):
-    image = cv2.imread(image_path)
-    if image is None:
-        return 0
+# YOLOv4-tiny (fast, decent accuracy, ~23MB weights)
+YOLO_CFG_URL = "https://raw.githubusercontent.com/AlexeyAB/darknet/master/cfg/yolov4-tiny.cfg"
+YOLO_WEIGHTS_URL = "https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v4_pre/yolov4-tiny.weights"
+YOLO_NAMES_URL = "https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names"
 
-    # Load OpenCV pre-trained model (HOG + SVM)
-    hog = cv2.HOGDescriptor()
-    hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
+MODEL_DIR = Path(os.getenv("MODEL_DIR", "models"))
+CFG_PATH = MODEL_DIR / "yolov4-tiny.cfg"
+WEIGHTS_PATH = MODEL_DIR / "yolov4-tiny.weights"
+NAMES_PATH = MODEL_DIR / "coco.names"
 
-    (rects, _) = hog.detectMultiScale(image, winStride=(4, 4), padding=(8, 8), scale=1.05)
-    return len(rects)
+def _ensure_models():
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    if not CFG_PATH.exists():
+        urllib.request.urlretrieve(YOLO_CFG_URL, CFG_PATH)
+    if not WEIGHTS_PATH.exists():
+        urllib.request.urlretrieve(YOLO_WEIGHTS_URL, WEIGHTS_PATH)
+    if not NAMES_PATH.exists():
+        urllib.request.urlretrieve(YOLO_NAMES_URL, NAMES_PATH)
+
+    with open(NAMES_PATH, "r") as f:
+        classes = [line.strip() for line in f.readlines()]
+    return classes
+
+_net = None
+_output_layers = None
+_classes = None
+
+def _load_net():
+    global _net, _output_layers, _classes
+    if _net is not None:
+        return _net, _output_layers, _classes
+    _classes = _ensure_models()
+    _net = cv2.dnn.readNetFromDarknet(str(CFG_PATH), str(WEIGHTS_PATH))
+    _net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
+    _net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
+    layer_names = _net.getLayerNames()
+    _output_layers = [layer_names[i - 1] for i in _net.getUnconnectedOutLayers().flatten()]
+    return _net, _output_layers, _classes
+
+def detect_people_yolo(frame_bgr, conf_thresh=0.55, nms_thresh=0.45, draw=True):
+    """
+    Returns:
+      people_indices: list of indices of 'person' boxes after NMS
+      boxes: list[(x,y,w,h)]
+      confidences: list[float]
+      annotated: frame with boxes (BGR)
+    """
+    net, out_layers, classes = _load_net()
+
+    h, w = frame_bgr.shape[:2]
+    blob = cv2.dnn.blobFromImage(frame_bgr, scalefactor=1/255.0, size=(416, 416),
+                                 swapRB=True, crop=False)
+    net.setInput(blob)
+    layer_outputs = net.forward(out_layers)
+
+    boxes = []
+    confidences = []
+    class_ids = []
+
+    for output in layer_outputs:
+        for detection in output:
+            scores = detection[5:]
+            class_id = int(np.argmax(scores))
+            confidence = float(scores[class_id])
+            if confidence < conf_thresh:
+                continue
+            center_x = int(detection[0] * w)
+            center_y = int(detection[1] * h)
+            bw = int(detection[2] * w)
+            bh = int(detection[3] * h)
+            x = int(center_x - bw / 2)
+            y = int(center_y - bh / 2)
+            boxes.append([x, y, bw, bh])
+            confidences.append(confidence)
+            class_ids.append(class_id)
+
+    idxs = cv2.dnn.NMSBoxes(boxes, confidences, conf_thresh, nms_thresh)
+
+    people_indices = []
+    annotated = frame_bgr.copy()
+    if len(idxs) > 0:
+        for i in idxs.flatten():
+            if class_ids[i] < len(classes) and classes[class_ids[i]] == "person":
+                people_indices.append(i)
+                if draw:
+                    x, y, bw, bh = boxes[i]
+                    cv2.rectangle(annotated, (x, y), (x + bw, y + bh), (0, 255, 0), 2)
+                    label = f"person {confidences[i]:.2f}"
+                    cv2.putText(annotated, label, (x, y - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+
+    # Big on-screen counter for debugging
+    if draw:
+        cv2.putText(annotated, f"People: {len(people_indices)}", (12, 36),
+                    cv2.FONT_HERSHEY_SIMPLEX, 1.1, (0, 0, 0), 4)
+        cv2.putText(annotated, f"People: {len(people_indices)}", (12, 36),
+                    cv2.FONT_HERSHEY_SIMPLEX, 1.1, (255, 255, 255), 2)
+
+    return people_indices, boxes, confidences, annotated