291 lines
11 KiB
Python
291 lines
11 KiB
Python
import cv2
|
|
import numpy as np
|
|
import yaml
|
|
import torch
|
|
from ultralytics import YOLO
|
|
import time
|
|
import datetime
|
|
import threading
|
|
import queue
|
|
import sys
|
|
import argparse
|
|
from collections import defaultdict
|
|
|
|
|
|
class ThreadedFrameReader:
|
|
def __init__(self, cam_id, rtsp_url):
|
|
self.cam_id = cam_id
|
|
self.rtsp_url = rtsp_url
|
|
self.cap = cv2.VideoCapture(rtsp_url, cv2.CAP_FFMPEG)
|
|
self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
|
self.q = queue.Queue(maxsize=2)
|
|
self.running = True
|
|
self.thread = threading.Thread(target=self._reader, 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 self.q.full():
|
|
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()
|
|
|
|
|
|
class MultiCameraMonitor:
|
|
def __init__(self, config_path):
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
self.cfg = yaml.safe_load(f)
|
|
|
|
# === 全局模型（只加载一次）===
|
|
model_cfg = self.cfg['model']
|
|
self.device = model_cfg.get('device', 'auto')
|
|
if self.device == 'auto' or not self.device:
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
print(f"🚀 全局加载模型到 {self.device}...")
|
|
self.model = YOLO(model_cfg['path'])
|
|
self.model.to(self.device)
|
|
self.use_half = (self.device == 'cuda')
|
|
if self.use_half:
|
|
print("✅ 启用 FP16 推理")
|
|
|
|
self.imgsz = model_cfg['imgsz']
|
|
self.conf_thresh = model_cfg['conf_threshold']
|
|
|
|
# === 初始化所有摄像头 ===
|
|
self.common = self.cfg['common']
|
|
self.cameras = {}
|
|
self.frame_readers = {}
|
|
self.queues = {} # cam_id -> queue for detection results
|
|
|
|
for cam_cfg in self.cfg['cameras']:
|
|
cam_id = cam_cfg['id']
|
|
self.cameras[cam_id] = CameraLogic(cam_id, cam_cfg, self.common)
|
|
self.frame_readers[cam_id] = ThreadedFrameReader(cam_id, cam_cfg['rtsp_url'])
|
|
self.queues[cam_id] = queue.Queue(maxsize=1) # 存放检测结果
|
|
|
|
# === 控制信号 ===
|
|
self.running = True
|
|
self.inference_thread = threading.Thread(target=self._inference_loop, daemon=True)
|
|
self.inference_thread.start()
|
|
|
|
def _inference_loop(self):
|
|
"""统一推理线程：轮询各摄像头最新帧，逐个推理"""
|
|
while self.running:
|
|
processed = False
|
|
for cam_id, reader in self.frame_readers.items():
|
|
ret, frame = reader.read()
|
|
if not ret:
|
|
continue
|
|
|
|
# 获取该摄像头是否需要处理此帧
|
|
cam_logic = self.cameras[cam_id]
|
|
if cam_logic.should_skip_frame():
|
|
continue
|
|
|
|
# 执行推理
|
|
results = self.model(
|
|
frame,
|
|
imgsz=self.imgsz,
|
|
conf=self.conf_thresh,
|
|
verbose=False,
|
|
device=self.device,
|
|
half=self.use_half,
|
|
classes=[0] # person
|
|
)
|
|
|
|
# 将结果送回对应摄像头逻辑
|
|
if not self.queues[cam_id].full():
|
|
self.queues[cam_id].put((frame.copy(), results[0]))
|
|
processed = True
|
|
|
|
if not processed:
|
|
time.sleep(0.01)
|
|
|
|
def run(self):
|
|
"""启动所有摄像头的显示和告警逻辑（主线程）"""
|
|
try:
|
|
while self.running:
|
|
for cam_id, cam_logic in self.cameras.items():
|
|
if not self.queues[cam_id].empty():
|
|
frame, results = self.queues[cam_id].get()
|
|
cam_logic.process(frame, results)
|
|
# 让 OpenCV 刷新所有窗口
|
|
key = cv2.waitKey(1) & 0xFF
|
|
if key == ord('q'): # 可选：按 q 退出
|
|
break
|
|
time.sleep(0.01) # 避免 CPU 占用过高
|
|
except KeyboardInterrupt:
|
|
pass
|
|
finally:
|
|
self.stop()
|
|
|
|
def stop(self):
|
|
self.running = False
|
|
for reader in self.frame_readers.values():
|
|
reader.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
class CameraLogic:
|
|
def __init__(self, cam_id, cam_cfg, common_cfg):
|
|
self.cam_id = cam_id
|
|
self.roi_off_duty = np.array(cam_cfg['roi_points'], dtype=np.int32)
|
|
self.roi_crowd = np.array(cam_cfg['crowd_roi_points'], dtype=np.int32)
|
|
self.off_duty_threshold_sec = cam_cfg['off_duty_threshold_sec']
|
|
self.on_duty_confirm_sec = cam_cfg['on_duty_confirm_sec']
|
|
self.off_duty_confirm_sec = cam_cfg['off_duty_confirm_sec']
|
|
self.crowd_threshold_min = cam_cfg['crowd_threshold']
|
|
|
|
self.working_hours = common_cfg['working_hours']
|
|
self.process_every_n = cam_cfg.get('process_every_n_frames', common_cfg['process_every_n_frames'])
|
|
self.alert_cooldown_sec = cam_cfg.get('alert_cooldown_sec', common_cfg['alert_cooldown_sec'])
|
|
self.crowd_cooldown_sec = cam_cfg.get('crowd_cooldown_sec', common_cfg['crowd_cooldown_sec'])
|
|
|
|
self.frame_count = 0
|
|
self.is_on_duty = False
|
|
self.is_off_duty = True
|
|
self.on_duty_start_time = None
|
|
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 = []
|
|
self.last_person_seen_time = None
|
|
|
|
def should_skip_frame(self):
|
|
self.frame_count += 1
|
|
return self.frame_count % self.process_every_n != 0
|
|
|
|
def is_point_in_roi(self, x, y, roi):
|
|
return cv2.pointPolygonTest(roi, (int(x), int(y)), False) >= 0
|
|
|
|
def in_working_hours(self):
|
|
now = datetime.datetime.now()
|
|
h = now.hour
|
|
start_h, end_h = self.working_hours
|
|
return start_h <= h < end_h or (start_h == end_h == 0)
|
|
|
|
def detect_crowd_count(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 self.is_point_in_roi(cx, cy, roi):
|
|
count += 1
|
|
return count
|
|
|
|
def process(self, frame, results):
|
|
current_time = time.time()
|
|
now = datetime.datetime.now()
|
|
in_work = self.in_working_hours()
|
|
boxes = results.boxes
|
|
|
|
# === 离岗逻辑 ===
|
|
roi_has_person_raw = any(
|
|
self.is_point_in_roi((b.xyxy[0][0] + b.xyxy[0][2]) / 2,
|
|
(b.xyxy[0][1] + b.xyxy[0][3]) / 2,
|
|
self.roi_off_duty)
|
|
for b in boxes
|
|
)
|
|
|
|
if in_work:
|
|
if roi_has_person_raw:
|
|
self.last_person_seen_time = current_time
|
|
effective = (
|
|
self.last_person_seen_time is not None and
|
|
(current_time - self.last_person_seen_time) < 1.0 # grace period
|
|
)
|
|
|
|
if effective:
|
|
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.on_duty_confirm_sec:
|
|
self.is_on_duty, self.is_off_duty = True, False
|
|
self.on_duty_start_time = None
|
|
print(f"[{self.cam_id}] ✅ 上岗确认成功 ({now.strftime('%H:%M:%S')})")
|
|
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.off_duty_confirm_sec:
|
|
self.is_off_duty, self.is_on_duty = True, False
|
|
self.last_no_person_time = None
|
|
self.off_duty_timer_start = current_time
|
|
print(f"[{self.cam_id}] 🚪 进入离岗计时")
|
|
|
|
# === 聚集检测 ===
|
|
crowd_count = self.detect_crowd_count(boxes, self.roi_crowd)
|
|
self.crowd_history.append((current_time, crowd_count))
|
|
self.crowd_history = [(t, c) for t, c in self.crowd_history if current_time - t <= 300]
|
|
|
|
if in_work and crowd_count >= self.crowd_threshold_min:
|
|
recent = [c for t, c in self.crowd_history[-10:]]
|
|
if len(recent) >= 5 and sum(c >= self.crowd_threshold_min for c in recent[-5:]) >= 4:
|
|
if current_time - self.last_crowd_alert_time >= self.crowd_cooldown_sec:
|
|
print(f"[{self.cam_id}] 🚨 聚集告警：{crowd_count}人")
|
|
self.last_crowd_alert_time = current_time
|
|
|
|
# === 离岗告警 ===
|
|
if in_work and self.is_off_duty and self.off_duty_timer_start:
|
|
if (current_time - self.off_duty_timer_start) >= self.off_duty_threshold_sec:
|
|
if (current_time - self.last_alert_time) >= self.alert_cooldown_sec:
|
|
print(f"[{self.cam_id}] 🚨 离岗告警！")
|
|
self.last_alert_time = current_time
|
|
|
|
# === 可视化 ===
|
|
vis = results.plot()
|
|
overlay = vis.copy()
|
|
cv2.fillPoly(overlay, [self.roi_off_duty], (0, 255, 0))
|
|
cv2.fillPoly(overlay, [self.roi_crowd], (0, 0, 255))
|
|
cv2.addWeighted(overlay, 0.2, vis, 0.8, 0, vis)
|
|
|
|
status = "OUT OF HOURS"
|
|
color = (128, 128, 128)
|
|
if in_work:
|
|
if self.is_on_duty:
|
|
status, color = "ON DUTY", (0, 255, 0)
|
|
elif self.is_off_duty:
|
|
if self.off_duty_timer_start:
|
|
elapsed = int(current_time - self.off_duty_timer_start)
|
|
if elapsed >= self.off_duty_threshold_sec:
|
|
status, color = "OFF DUTY!", (0, 0, 255)
|
|
else:
|
|
status = f"IDLE - {elapsed}s"
|
|
color = (0, 255, 255)
|
|
else:
|
|
status, color = "OFF DUTY", (255, 0, 0)
|
|
|
|
cv2.putText(vis, f"[{self.cam_id}] {status}", (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
|
|
cv2.imshow(f"Monitor - {self.cam_id}", vis)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default="config.yaml", help="配置文件路径")
|
|
args = parser.parse_args()
|
|
|
|
monitor = MultiCameraMonitor(args.config)
|
|
monitor.run()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |