detect_people.py

import os
import cv2
import numpy as np
from pathlib import Path
import urllib.request

# 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"

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"

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