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() |