218 lines
8.4 KiB
Python
218 lines
8.4 KiB
Python
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
from sort import Sort
|
|||
|
|
import time
|
|||
|
|
import datetime
|
|||
|
|
import threading
|
|||
|
|
import queue
|
|||
|
|
import torch
|
|||
|
|
from collections import deque
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ThreadedFrameReader:
|
|||
|
|
def __init__(self, src, maxsize=1):
|
|||
|
|
self.cap = cv2.VideoCapture(src, cv2.CAP_FFMPEG)
|
|||
|
|
self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
|||
|
|
self.q = queue.Queue(maxsize=maxsize)
|
|||
|
|
self.running = True
|
|||
|
|
self.thread = threading.Thread(target=self._reader)
|
|||
|
|
self.thread.daemon = True
|
|||
|
|
self.thread.start()
|
|||
|
|
|
|||
|
|
def _reader(self):
|
|||
|
|
while self.running:
|
|||
|
|
ret, frame = self.cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
time.sleep(0.1)
|
|||
|
|
continue
|
|||
|
|
if not self.q.empty():
|
|||
|
|
try:
|
|||
|
|
self.q.get_nowait()
|
|||
|
|
except queue.Empty:
|
|||
|
|
pass
|
|||
|
|
self.q.put(frame)
|
|||
|
|
|
|||
|
|
def read(self):
|
|||
|
|
if not self.q.empty():
|
|||
|
|
return True, self.q.get()
|
|||
|
|
return False, None
|
|||
|
|
|
|||
|
|
def release(self):
|
|||
|
|
self.running = False
|
|||
|
|
self.cap.release()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_point_in_roi(x, y, roi):
|
|||
|
|
return cv2.pointPolygonTest(roi, (int(x), int(y)), False) >= 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OffDutyCrowdDetector:
|
|||
|
|
def __init__(self, config, model, device, use_half):
|
|||
|
|
self.config = config
|
|||
|
|
self.model = model
|
|||
|
|
self.device = device
|
|||
|
|
self.use_half = use_half
|
|||
|
|
|
|||
|
|
# 解析 ROI
|
|||
|
|
self.roi = np.array(config["roi_points"], dtype=np.int32)
|
|||
|
|
self.crowd_roi = np.array(config["crowd_roi_points"], dtype=np.int32)
|
|||
|
|
|
|||
|
|
# 状态变量
|
|||
|
|
self.tracker = Sort(
|
|||
|
|
max_age=30,
|
|||
|
|
min_hits=2,
|
|||
|
|
iou_threshold=0.3
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self.is_on_duty = False
|
|||
|
|
self.on_duty_start_time = None
|
|||
|
|
self.is_off_duty = True
|
|||
|
|
self.last_no_person_time = None
|
|||
|
|
self.off_duty_timer_start = None
|
|||
|
|
self.last_alert_time = 0
|
|||
|
|
|
|||
|
|
self.last_crowd_alert_time = 0
|
|||
|
|
self.crowd_history = deque(maxlen=1500) # 自动限制5分钟（假设5fps）
|
|||
|
|
|
|||
|
|
self.last_person_seen_time = None
|
|||
|
|
self.frame_count = 0
|
|||
|
|
|
|||
|
|
# 缓存配置
|
|||
|
|
self.working_start_min = config.get("working_hours", [9, 17])[0] * 60
|
|||
|
|
self.working_end_min = config.get("working_hours", [9, 17])[1] * 60
|
|||
|
|
self.process_every = config.get("process_every_n_frames", 3)
|
|||
|
|
|
|||
|
|
def in_working_hours(self):
|
|||
|
|
now = datetime.datetime.now()
|
|||
|
|
total_min = now.hour * 60 + now.minute
|
|||
|
|
return self.working_start_min <= total_min <= self.working_end_min
|
|||
|
|
|
|||
|
|
def count_people_in_roi(self, boxes, roi):
|
|||
|
|
count = 0
|
|||
|
|
for box in boxes:
|
|||
|
|
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
|
|||
|
|
cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
|
|||
|
|
if is_point_in_roi(cx, cy, roi):
|
|||
|
|
count += 1
|
|||
|
|
return count
|
|||
|
|
|
|||
|
|
def run(self):
|
|||
|
|
"""主循环：供线程调用"""
|
|||
|
|
frame_reader = ThreadedFrameReader(self.config["rtsp_url"])
|
|||
|
|
try:
|
|||
|
|
while True:
|
|||
|
|
ret, frame = frame_reader.read()
|
|||
|
|
if not ret:
|
|||
|
|
time.sleep(0.01)
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
self.frame_count += 1
|
|||
|
|
if self.frame_count % self.process_every != 0:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
current_time = time.time()
|
|||
|
|
now = datetime.datetime.now()
|
|||
|
|
|
|||
|
|
# YOLO 推理
|
|||
|
|
results = self.model(
|
|||
|
|
frame,
|
|||
|
|
imgsz=self.config.get("imgsz", 480),
|
|||
|
|
conf=self.config.get("conf_thresh", 0.4),
|
|||
|
|
verbose=False,
|
|||
|
|
device=self.device,
|
|||
|
|
half=self.use_half,
|
|||
|
|
classes=[0] # person class
|
|||
|
|
)
|
|||
|
|
boxes = results[0].boxes
|
|||
|
|
|
|||
|
|
# 更新 tracker（可选，用于ID跟踪）
|
|||
|
|
dets = []
|
|||
|
|
for box in boxes:
|
|||
|
|
x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
|
|||
|
|
conf = float(box.conf)
|
|||
|
|
dets.append([x1, y1, x2, y2, conf])
|
|||
|
|
dets = np.array(dets) if dets else np.empty((0, 5))
|
|||
|
|
self.tracker.update(dets)
|
|||
|
|
|
|||
|
|
# === 离岗检测 ===
|
|||
|
|
if self.in_working_hours():
|
|||
|
|
roi_has_person = self.count_people_in_roi(boxes, self.roi) > 0
|
|||
|
|
if roi_has_person:
|
|||
|
|
self.last_person_seen_time = current_time
|
|||
|
|
|
|||
|
|
# 入岗保护期
|
|||
|
|
effective_on_duty = (
|
|||
|
|
self.last_person_seen_time is not None and
|
|||
|
|
(current_time - self.last_person_seen_time) < 1.0
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if effective_on_duty:
|
|||
|
|
self.last_no_person_time = None
|
|||
|
|
if self.is_off_duty:
|
|||
|
|
if self.on_duty_start_time is None:
|
|||
|
|
self.on_duty_start_time = current_time
|
|||
|
|
elif current_time - self.on_duty_start_time >= self.config.get("on_duty_confirm", 5):
|
|||
|
|
self.is_on_duty = True
|
|||
|
|
self.is_off_duty = False
|
|||
|
|
self.on_duty_start_time = None
|
|||
|
|
print(f"[{self.config['id']}] ✅ 上岗确认")
|
|||
|
|
else:
|
|||
|
|
self.on_duty_start_time = None
|
|||
|
|
self.last_person_seen_time = None
|
|||
|
|
if not self.is_off_duty:
|
|||
|
|
if self.last_no_person_time is None:
|
|||
|
|
self.last_no_person_time = current_time
|
|||
|
|
elif current_time - self.last_no_person_time >= self.config.get("off_duty_confirm", 30):
|
|||
|
|
self.is_off_duty = True
|
|||
|
|
self.is_on_duty = False
|
|||
|
|
self.off_duty_timer_start = current_time
|
|||
|
|
print(f"[{self.config['id']}] ⏳ 开始离岗计时")
|
|||
|
|
|
|||
|
|
# 离岗告警
|
|||
|
|
if self.is_off_duty and self.off_duty_timer_start:
|
|||
|
|
elapsed = current_time - self.off_duty_timer_start
|
|||
|
|
if elapsed >= self.config.get("off_duty_threshold", 300):
|
|||
|
|
if current_time - self.last_alert_time >= self.config.get("alert_cooldown", 300):
|
|||
|
|
print(f"[{self.config['id']}] 🚨 离岗告警！已离岗 {elapsed/60:.1f} 分钟")
|
|||
|
|
self.last_alert_time = current_time
|
|||
|
|
|
|||
|
|
# === 聚集检测 ===
|
|||
|
|
crowd_count = self.count_people_in_roi(boxes, self.crowd_roi)
|
|||
|
|
self.crowd_history.append((current_time, crowd_count))
|
|||
|
|
|
|||
|
|
# 动态阈值
|
|||
|
|
if crowd_count >= 10:
|
|||
|
|
req_dur = 60
|
|||
|
|
elif crowd_count >= 7:
|
|||
|
|
req_dur = 120
|
|||
|
|
elif crowd_count >= 5:
|
|||
|
|
req_dur = 300
|
|||
|
|
else:
|
|||
|
|
req_dur = float('inf')
|
|||
|
|
|
|||
|
|
if req_dur < float('inf'):
|
|||
|
|
recent = [(t, c) for t, c in self.crowd_history if current_time - t <= req_dur]
|
|||
|
|
if recent:
|
|||
|
|
valid = [c for t, c in recent if c >= 4]
|
|||
|
|
ratio = len(valid) / len(recent)
|
|||
|
|
if ratio >= 0.9 and (current_time - self.last_crowd_alert_time) >= self.config.get("crowd_cooldown", 180):
|
|||
|
|
print(f"[{self.config['id']}] 🚨 聚集告警！{crowd_count}人持续{req_dur//60}分钟")
|
|||
|
|
self.last_crowd_alert_time = current_time
|
|||
|
|
|
|||
|
|
# 可视化（可选，部署时可关闭）
|
|||
|
|
if True: # 设为 True 可显示窗口
|
|||
|
|
vis = results[0].plot()
|
|||
|
|
overlay = vis.copy()
|
|||
|
|
cv2.fillPoly(overlay, [self.roi], (0,255,0))
|
|||
|
|
cv2.fillPoly(overlay, [self.crowd_roi], (0,0,255))
|
|||
|
|
cv2.addWeighted(overlay, 0.2, vis, 0.8, 0, vis)
|
|||
|
|
cv2.imshow(f"Monitor - {self.config['id']}", vis)
|
|||
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"[{self.config['id']}] Error: {e}")
|
|||
|
|
finally:
|
|||
|
|
frame_reader.release()
|
|||
|
|
cv2.destroyAllWindows()
|