Files
security-ai-edge/algorithms.py
16337 c4baf2fd1f fix: 重构算法逻辑
LeavePostAlgorithm:
- 重构状态机:WAITING  ON_DUTY  LEAVING  OFF_DUTY
- 移除 INIT 状态,增加 LEAVING 确认期防抖动
- 添加目标类型过滤 (target_class)
- 统一使用 datetime

IntrusionAlgorithm:
- 修复冷却逻辑 Bug:只在告警后重置
- 添加确认机制:持续检测 N 秒才触发告警
- 添加目标类型过滤 (target_class)
- 二次告警时间改为 2 分钟 (120秒)

AlgorithmManager:
- 更新默认参数
- 统一参数命名
2026-01-30 13:23:22 +08:00

374 lines
13 KiB
Python

import os
import sys
import time
from collections import deque
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Tuple
import cv2
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
class LeavePostAlgorithm:
STATE_WAITING = "WAITING"
STATE_ON_DUTY = "ON_DUTY"
STATE_LEAVING = "LEAVING"
STATE_OFF_DUTY = "OFF_DUTY"
STATE_NON_WORK_TIME = "NON_WORK_TIME"
def __init__(
self,
confirm_on_duty_sec: int = 10,
confirm_leave_sec: int = 10,
cooldown_sec: int = 300,
working_hours: Optional[List[Dict]] = None,
target_class: Optional[str] = "person",
):
self.confirm_on_duty_sec = confirm_on_duty_sec
self.confirm_leave_sec = confirm_leave_sec
self.cooldown_sec = cooldown_sec
self.working_hours = working_hours or []
self.target_class = target_class
self.alert_cooldowns: Dict[str, datetime] = {}
self.state: str = self.STATE_WAITING
self.state_start_time: Optional[datetime] = None
self.detection_history: deque = deque()
self.alarm_sent: bool = False
self.last_person_time: Optional[datetime] = None
def _is_in_working_hours(self, dt: Optional[datetime] = None) -> bool:
if not self.working_hours:
return True
dt = dt or datetime.now()
current_minutes = dt.hour * 60 + dt.minute
for period in self.working_hours:
start_minutes = period["start"][0] * 60 + period["start"][1]
end_minutes = period["end"][0] * 60 + period["end"][1]
if start_minutes <= current_minutes < end_minutes:
return True
return False
def _check_detection_in_roi(self, detection: Dict, roi_id: str) -> bool:
matched_rois = detection.get("matched_rois", [])
for roi in matched_rois:
if roi.get("roi_id") == roi_id:
return True
return False
def _check_target_class(self, detection: Dict, target_class: Optional[str]) -> bool:
if not target_class:
return True
return detection.get("class") == target_class
def _get_detection_window(self, current_time: datetime) -> List[bool]:
detections = []
while self.detection_history and (current_time - self.detection_history[0][0]).total_seconds() > max(self.confirm_on_duty_sec, self.confirm_leave_sec):
self.detection_history.popleft()
for _, has_person in self.detection_history:
detections.append(has_person)
return detections
def _get_latest_bbox(self, tracks: List[Dict], roi_id: str) -> List[float]:
for det in tracks:
if self._check_detection_in_roi(det, roi_id):
return det.get("bbox", [])
return []
def process(
self,
roi_id: str,
camera_id: str,
tracks: List[Dict],
current_time: Optional[datetime] = None,
) -> List[Dict]:
current_time = current_time or datetime.now()
in_work = self._is_in_working_hours(current_time)
alerts = []
if not in_work:
self.state = self.STATE_NON_WORK_TIME
self.detection_history.clear()
self.alarm_sent = False
return []
if self.state == self.STATE_NON_WORK_TIME:
self.state = self.STATE_WAITING
self.detection_history.clear()
self.alarm_sent = False
roi_has_person = False
for det in tracks:
if self._check_detection_in_roi(det, roi_id) and self._check_target_class(det, self.target_class):
roi_has_person = True
break
if self.state == self.STATE_WAITING:
if roi_has_person:
self.state = self.STATE_ON_DUTY
self.state_start_time = current_time
self.detection_history.clear()
self.detection_history.append((current_time, True))
else:
pass
elif self.state == self.STATE_ON_DUTY:
self.detection_history.append((current_time, roi_has_person))
if not roi_has_person:
self.state = self.STATE_LEAVING
self.state_start_time = current_time
elif self.state == self.STATE_LEAVING:
self.detection_history.append((current_time, roi_has_person))
elapsed = (current_time - self.state_start_time).total_seconds()
if roi_has_person:
self.state = self.STATE_ON_DUTY
self.state_start_time = current_time
elif elapsed >= self.confirm_leave_sec:
self.state = self.STATE_OFF_DUTY
self.state_start_time = current_time
elif self.state == self.STATE_OFF_DUTY:
elapsed = (current_time - self.state_start_time).total_seconds()
if roi_has_person:
self.state = self.STATE_ON_DUTY
self.state_start_time = current_time
elif elapsed >= self.confirm_leave_sec:
cooldown_key = f"{camera_id}_{roi_id}"
now = datetime.now()
if cooldown_key not in self.alert_cooldowns or (now - self.alert_cooldowns[cooldown_key]).total_seconds() > self.cooldown_sec:
bbox = self._get_latest_bbox(tracks, roi_id)
elapsed_minutes = int(elapsed / 60)
alerts.append({
"track_id": roi_id,
"camera_id": camera_id,
"bbox": bbox,
"off_duty_duration": elapsed,
"alert_type": "leave_post",
"message": f"离岗超过 {elapsed_minutes} 分钟",
})
self.alert_cooldowns[cooldown_key] = now
return alerts
def reset(self):
self.state = self.STATE_WAITING
self.state_start_time = None
self.detection_history.clear()
self.alarm_sent = False
self.last_person_time = None
self.alert_cooldowns.clear()
def get_state(self, roi_id: str) -> Dict[str, Any]:
return {
"state": self.state,
"alarm_sent": self.alarm_sent,
"last_person_time": self.last_person_time,
}
class IntrusionAlgorithm:
def __init__(
self,
cooldown_seconds: int = 120,
confirm_seconds: int = 3,
target_class: Optional[str] = None,
):
self.cooldown_seconds = cooldown_seconds
self.confirm_seconds = confirm_seconds
self.target_class = target_class
self.last_alert_time: Dict[str, datetime] = {}
self.alert_triggered: Dict[str, bool] = {}
self.detection_start: Dict[str, Optional[datetime]] = {}
def _check_detection_in_roi(self, detection: Dict, roi_id: str) -> bool:
matched_rois = detection.get("matched_rois", [])
for roi in matched_rois:
if roi.get("roi_id") == roi_id:
return True
return False
def _check_target_class(self, detection: Dict, target_class: Optional[str]) -> bool:
if not target_class:
return True
return detection.get("class") == target_class
def _get_latest_bbox(self, tracks: List[Dict], roi_id: str) -> List[float]:
for det in tracks:
if self._check_detection_in_roi(det, roi_id):
return det.get("bbox", [])
return []
def process(
self,
roi_id: str,
camera_id: str,
tracks: List[Dict],
current_time: Optional[datetime] = None,
) -> List[Dict]:
current_time = current_time or datetime.now()
key = f"{camera_id}_{roi_id}"
roi_has_person = False
for det in tracks:
if self._check_detection_in_roi(det, roi_id) and self._check_target_class(det, self.target_class):
roi_has_person = True
break
if not roi_has_person:
self.detection_start.pop(key, None)
self.alert_triggered[key] = False
return []
if self.alert_triggered.get(key, False):
elapsed_since_alert = (current_time - self.last_alert_time.get(key, datetime.min)).total_seconds()
if elapsed_since_alert < self.cooldown_seconds:
return []
self.alert_triggered[key] = False
if self.detection_start.get(key) is None:
self.detection_start[key] = current_time
elapsed = (current_time - self.detection_start[key]).total_seconds()
if elapsed < self.confirm_seconds:
return []
bbox = self._get_latest_bbox(tracks, roi_id)
self.last_alert_time[key] = current_time
self.alert_triggered[key] = True
self.detection_start[key] = None
return [{
"roi_id": roi_id,
"camera_id": camera_id,
"bbox": bbox,
"alert_type": "intrusion",
"message": "检测到周界入侵",
}]
def reset(self):
self.last_alert_time.clear()
self.alert_triggered.clear()
self.detection_start.clear()
class AlgorithmManager:
def __init__(self, working_hours: Optional[List[Dict]] = None):
self.algorithms: Dict[str, Dict[str, Any]] = {}
self.working_hours = working_hours or []
self.default_params = {
"leave_post": {
"confirm_on_duty_sec": 10,
"confirm_leave_sec": 10,
"cooldown_sec": 300,
"target_class": "person",
},
"intrusion": {
"cooldown_seconds": 120,
"confirm_seconds": 3,
"target_class": None,
},
}
def register_algorithm(
self,
roi_id: str,
algorithm_type: str,
params: Optional[Dict[str, Any]] = None,
):
if roi_id in self.algorithms and algorithm_type in self.algorithms[roi_id]:
return
if roi_id not in self.algorithms:
self.algorithms[roi_id] = {}
algo_params = self.default_params.get(algorithm_type, {}).copy()
if params:
algo_params.update(params)
if algorithm_type == "leave_post":
roi_working_hours = algo_params.get("working_hours") or self.working_hours
self.algorithms[roi_id]["leave_post"] = LeavePostAlgorithm(
confirm_on_duty_sec=algo_params.get("confirm_on_duty_sec", 10),
confirm_leave_sec=algo_params.get("confirm_leave_sec", 10),
cooldown_sec=algo_params.get("cooldown_sec", 300),
working_hours=roi_working_hours,
target_class=algo_params.get("target_class", "person"),
)
elif algorithm_type == "intrusion":
self.algorithms[roi_id]["intrusion"] = IntrusionAlgorithm(
cooldown_seconds=algo_params.get("cooldown_seconds", 120),
confirm_seconds=algo_params.get("confirm_seconds", 3),
target_class=algo_params.get("target_class"),
)
def process(
self,
roi_id: str,
camera_id: str,
algorithm_type: str,
tracks: List[Dict],
current_time: Optional[datetime] = None,
) -> List[Dict]:
algo = self.algorithms.get(roi_id, {}).get(algorithm_type)
if algo is None:
return []
return algo.process(roi_id, camera_id, tracks, current_time)
def update_roi_params(
self,
roi_id: str,
algorithm_type: str,
params: Dict[str, Any],
):
if roi_id in self.algorithms and algorithm_type in self.algorithms[roi_id]:
algo = self.algorithms[roi_id][algorithm_type]
for key, value in params.items():
if hasattr(algo, key):
setattr(algo, key, value)
def reset_algorithm(self, roi_id: str, algorithm_type: Optional[str] = None):
if roi_id not in self.algorithms:
return
if algorithm_type:
if algorithm_type in self.algorithms[roi_id]:
self.algorithms[roi_id][algorithm_type].reset()
else:
for algo in self.algorithms[roi_id].values():
algo.reset()
def reset_all(self):
for roi_algorithms in self.algorithms.values():
for algo in roi_algorithms.values():
algo.reset()
def remove_roi(self, roi_id: str):
if roi_id in self.algorithms:
self.reset_algorithm(roi_id)
del self.algorithms[roi_id]
def get_status(self, roi_id: str) -> Dict[str, Any]:
status = {}
if roi_id in self.algorithms:
for algo_type, algo in self.algorithms[roi_id].items():
if algo_type == "leave_post":
status[algo_type] = {
"state": getattr(algo, "state", "WAITING"),
"alarm_sent": getattr(algo, "alarm_sent", False),
}
else:
status[algo_type] = {
"detection_count": len(getattr(algo, "detection_start", {})),
}
return status