304 lines
10 KiB
Python
304 lines
10 KiB
Python
|
import os
|
||
|
|
import sys
|
||
|
|
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__))))
|
||
|
|
|
||
|
|
from sort import Sort
|
||
|
|
|
||
|
|
|
||
|
|
class LeavePostAlgorithm:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
threshold_sec: int = 360,
|
||
|
|
confirm_sec: int = 30,
|
||
|
|
return_sec: int = 5,
|
||
|
|
working_hours: Optional[List[Dict]] = None,
|
||
|
|
):
|
||
|
|
self.threshold_sec = threshold_sec
|
||
|
|
self.confirm_sec = confirm_sec
|
||
|
|
self.return_sec = return_sec
|
||
|
|
self.working_hours = working_hours or []
|
||
|
|
|
||
|
|
self.track_states: Dict[str, Dict[str, Any]] = {}
|
||
|
|
self.tracker = Sort(max_age=10, min_hits=2, iou_threshold=0.3)
|
||
|
|
|
||
|
|
self.alert_cooldowns: Dict[str, datetime] = {}
|
||
|
|
self.cooldown_seconds = 300
|
||
|
|
|
||
|
|
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 process(
|
||
|
|
self,
|
||
|
|
camera_id: str,
|
||
|
|
tracks: List[Dict],
|
||
|
|
current_time: Optional[datetime] = None,
|
||
|
|
) -> List[Dict]:
|
||
|
|
if not self.is_in_working_hours(current_time):
|
||
|
|
return []
|
||
|
|
|
||
|
|
if not tracks:
|
||
|
|
return []
|
||
|
|
|
||
|
|
detections = []
|
||
|
|
for track in tracks:
|
||
|
|
bbox = track.get("bbox", [])
|
||
|
|
if len(bbox) >= 4:
|
||
|
|
detections.append(bbox + [track.get("conf", 0.0)])
|
||
|
|
|
||
|
|
if not detections:
|
||
|
|
return []
|
||
|
|
|
||
|
|
detections = np.array(detections)
|
||
|
|
tracked = self.tracker.update(detections)
|
||
|
|
|
||
|
|
alerts = []
|
||
|
|
current_time = current_time or datetime.now()
|
||
|
|
|
||
|
|
for track_data in tracked:
|
||
|
|
x1, y1, x2, y2, track_id = track_data
|
||
|
|
track_id = str(int(track_id))
|
||
|
|
|
||
|
|
if track_id not in self.track_states:
|
||
|
|
self.track_states[track_id] = {
|
||
|
|
"first_seen": current_time,
|
||
|
|
"last_seen": current_time,
|
||
|
|
"off_duty_start": None,
|
||
|
|
"alerted": False,
|
||
|
|
"last_position": (x1, y1, x2, y2),
|
||
|
|
}
|
||
|
|
|
||
|
|
state = self.track_states[track_id]
|
||
|
|
state["last_seen"] = current_time
|
||
|
|
state["last_position"] = (x1, y1, x2, y2)
|
||
|
|
|
||
|
|
if state["off_duty_start"] is None:
|
||
|
|
off_duty_duration = (current_time - state["first_seen"]).total_seconds()
|
||
|
|
if off_duty_duration > self.confirm_sec:
|
||
|
|
state["off_duty_start"] = current_time
|
||
|
|
else:
|
||
|
|
elapsed = (current_time - state["off_duty_start"]).total_seconds()
|
||
|
|
if elapsed > self.threshold_sec:
|
||
|
|
if not state["alerted"]:
|
||
|
|
cooldown_key = f"{camera_id}_{track_id}"
|
||
|
|
now = datetime.now()
|
||
|
|
if cooldown_key not in self.alert_cooldowns or (
|
||
|
|
now - self.alert_cooldowns[cooldown_key]
|
||
|
|
).total_seconds() > self.cooldown_seconds:
|
||
|
|
alerts.append({
|
||
|
|
"track_id": track_id,
|
||
|
|
"bbox": [x1, y1, x2, y2],
|
||
|
|
"off_duty_duration": elapsed,
|
||
|
|
"alert_type": "leave_post",
|
||
|
|
"message": f"离岗超过 {int(elapsed / 60)} 分钟",
|
||
|
|
})
|
||
|
|
state["alerted"] = True
|
||
|
|
self.alert_cooldowns[cooldown_key] = now
|
||
|
|
else:
|
||
|
|
if elapsed < self.return_sec:
|
||
|
|
state["off_duty_start"] = None
|
||
|
|
state["alerted"] = False
|
||
|
|
|
||
|
|
cleanup_time = current_time - timedelta(minutes=5)
|
||
|
|
for track_id, state in list(self.track_states.items()):
|
||
|
|
if state["last_seen"] < cleanup_time:
|
||
|
|
del self.track_states[track_id]
|
||
|
|
|
||
|
|
return alerts
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
self.track_states.clear()
|
||
|
|
self.alert_cooldowns.clear()
|
||
|
|
|
||
|
|
def get_state(self, track_id: str) -> Optional[Dict[str, Any]]:
|
||
|
|
return self.track_states.get(track_id)
|
||
|
|
|
||
|
|
|
||
|
|
class IntrusionAlgorithm:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
check_interval_sec: float = 1.0,
|
||
|
|
direction_sensitive: bool = False,
|
||
|
|
):
|
||
|
|
self.check_interval_sec = check_interval_sec
|
||
|
|
self.direction_sensitive = direction_sensitive
|
||
|
|
|
||
|
|
self.last_check_times: Dict[str, float] = {}
|
||
|
|
self.tracker = Sort(max_age=5, min_hits=1, iou_threshold=0.3)
|
||
|
|
|
||
|
|
self.alert_cooldowns: Dict[str, datetime] = {}
|
||
|
|
self.cooldown_seconds = 300
|
||
|
|
|
||
|
|
def process(
|
||
|
|
self,
|
||
|
|
camera_id: str,
|
||
|
|
tracks: List[Dict],
|
||
|
|
current_time: Optional[datetime] = None,
|
||
|
|
) -> List[Dict]:
|
||
|
|
if not tracks:
|
||
|
|
return []
|
||
|
|
|
||
|
|
detections = []
|
||
|
|
for track in tracks:
|
||
|
|
bbox = track.get("bbox", [])
|
||
|
|
if len(bbox) >= 4:
|
||
|
|
detections.append(bbox + [track.get("conf", 0.0)])
|
||
|
|
|
||
|
|
if not detections:
|
||
|
|
return []
|
||
|
|
|
||
|
|
current_ts = current_time.timestamp() if current_time else datetime.now().timestamp()
|
||
|
|
|
||
|
|
if camera_id in self.last_check_times:
|
||
|
|
if current_ts - self.last_check_times[camera_id] < self.check_interval_sec:
|
||
|
|
return []
|
||
|
|
self.last_check_times[camera_id] = current_ts
|
||
|
|
|
||
|
|
detections = np.array(detections)
|
||
|
|
tracked = self.tracker.update(detections)
|
||
|
|
|
||
|
|
alerts = []
|
||
|
|
now = datetime.now()
|
||
|
|
|
||
|
|
for track_data in tracked:
|
||
|
|
x1, y1, x2, y2, track_id = track_data
|
||
|
|
cooldown_key = f"{camera_id}_{int(track_id)}"
|
||
|
|
|
||
|
|
if cooldown_key not in self.alert_cooldowns or (
|
||
|
|
now - self.alert_cooldowns[cooldown_key]
|
||
|
|
).total_seconds() > self.cooldown_seconds:
|
||
|
|
alerts.append({
|
||
|
|
"track_id": str(int(track_id)),
|
||
|
|
"bbox": [x1, y1, x2, y2],
|
||
|
|
"alert_type": "intrusion",
|
||
|
|
"confidence": track_data[4] if len(track_data) > 4 else 0.0,
|
||
|
|
"message": "检测到周界入侵",
|
||
|
|
})
|
||
|
|
self.alert_cooldowns[cooldown_key] = now
|
||
|
|
|
||
|
|
return alerts
|
||
|
|
|
||
|
|
def reset(self):
|
||
|
|
self.last_check_times.clear()
|
||
|
|
self.alert_cooldowns.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": {
|
||
|
|
"threshold_sec": 360,
|
||
|
|
"confirm_sec": 30,
|
||
|
|
"return_sec": 5,
|
||
|
|
},
|
||
|
|
"intrusion": {
|
||
|
|
"check_interval_sec": 1.0,
|
||
|
|
"direction_sensitive": False,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
def register_algorithm(
|
||
|
|
self,
|
||
|
|
roi_id: str,
|
||
|
|
algorithm_type: str,
|
||
|
|
params: Optional[Dict[str, Any]] = None,
|
||
|
|
):
|
||
|
|
if roi_id in self.algorithms:
|
||
|
|
if 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, {})
|
||
|
|
if params:
|
||
|
|
algo_params.update(params)
|
||
|
|
|
||
|
|
if algorithm_type == "leave_post":
|
||
|
|
self.algorithms[roi_id]["leave_post"] = LeavePostAlgorithm(
|
||
|
|
threshold_sec=algo_params.get("threshold_sec", 360),
|
||
|
|
confirm_sec=algo_params.get("confirm_sec", 30),
|
||
|
|
return_sec=algo_params.get("return_sec", 5),
|
||
|
|
working_hours=self.working_hours,
|
||
|
|
)
|
||
|
|
elif algorithm_type == "intrusion":
|
||
|
|
self.algorithms[roi_id]["intrusion"] = IntrusionAlgorithm(
|
||
|
|
check_interval_sec=algo_params.get("check_interval_sec", 1.0),
|
||
|
|
direction_sensitive=algo_params.get("direction_sensitive", False),
|
||
|
|
)
|
||
|
|
|
||
|
|
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(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():
|
||
|
|
status[algo_type] = {
|
||
|
|
"track_count": len(getattr(algo, "track_states", {})),
|
||
|
|
}
|
||
|
|
return status
|