Spaces:
Sleeping
Sleeping
File size: 3,948 Bytes
fc33912 7c4cf26 fc33912 7c4cf26 fc33912 c4ce88c fc33912 c4ce88c fc33912 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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
|