feat: add video detection and robot alerts
This commit is contained in:
323
backend/alerting.py
Normal file
323
backend/alerting.py
Normal file
@@ -0,0 +1,323 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
ALERT_CLASSES = ("fire", "smoke")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SessionState:
|
||||
consecutive: dict[str, int] = field(
|
||||
default_factory=lambda: defaultdict(int)
|
||||
)
|
||||
last_alert_at: dict[str, float] = field(
|
||||
default_factory=lambda: defaultdict(float)
|
||||
)
|
||||
last_seen_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
class AlertManager:
|
||||
def __init__(
|
||||
self,
|
||||
webhook_url: str | None = None,
|
||||
wechat_webhook_url: str | None = None,
|
||||
feishu_webhook_url: str | None = None,
|
||||
feishu_secret: str | None = None,
|
||||
confirm_frames: int = 3,
|
||||
cooldown_seconds: float = 60.0,
|
||||
session_ttl_seconds: float = 3600.0,
|
||||
) -> None:
|
||||
self.wechat_webhook_url = (
|
||||
wechat_webhook_url
|
||||
or webhook_url
|
||||
or os.getenv("WECHAT_WEBHOOK_URL")
|
||||
)
|
||||
self.feishu_webhook_url = (
|
||||
feishu_webhook_url or os.getenv("FEISHU_WEBHOOK_URL")
|
||||
)
|
||||
self.feishu_secret = feishu_secret or os.getenv("FEISHU_SECRET")
|
||||
self.confirm_frames = max(1, confirm_frames)
|
||||
self.cooldown_seconds = max(0.0, cooldown_seconds)
|
||||
self.session_ttl_seconds = max(60.0, session_ttl_seconds)
|
||||
self._states: dict[str, SessionState] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._executor = ThreadPoolExecutor(
|
||||
max_workers=4,
|
||||
thread_name_prefix="alert-dispatch",
|
||||
)
|
||||
|
||||
@property
|
||||
def channels(self) -> dict[str, bool]:
|
||||
return {
|
||||
"wechat": bool(self.wechat_webhook_url),
|
||||
"feishu": bool(self.feishu_webhook_url),
|
||||
}
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return any(self.channels.values())
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
session_id: str,
|
||||
detections: list[dict[str, Any]],
|
||||
image: Image.Image,
|
||||
) -> dict[str, Any]:
|
||||
now = time.monotonic()
|
||||
present = {
|
||||
detection["class"]
|
||||
for detection in detections
|
||||
if detection.get("class") in ALERT_CLASSES
|
||||
}
|
||||
triggered: list[str] = []
|
||||
|
||||
with self._lock:
|
||||
self._prune_sessions(now)
|
||||
state = self._states.setdefault(session_id, SessionState())
|
||||
state.last_seen_at = now
|
||||
for class_name in ALERT_CLASSES:
|
||||
state.consecutive[class_name] = (
|
||||
state.consecutive[class_name] + 1
|
||||
if class_name in present
|
||||
else 0
|
||||
)
|
||||
ready = state.consecutive[class_name] >= self.confirm_frames
|
||||
cooldown_elapsed = (
|
||||
now - state.last_alert_at[class_name]
|
||||
>= self.cooldown_seconds
|
||||
)
|
||||
if ready and cooldown_elapsed:
|
||||
state.last_alert_at[class_name] = now
|
||||
triggered.append(class_name)
|
||||
consecutive = dict(state.consecutive)
|
||||
|
||||
if triggered and self.enabled:
|
||||
self._executor.submit(
|
||||
self._send_alerts,
|
||||
annotate_image(image, detections),
|
||||
triggered,
|
||||
detections,
|
||||
)
|
||||
return {
|
||||
"triggered": bool(triggered),
|
||||
"classes": triggered,
|
||||
"confirmed_frames": self.confirm_frames,
|
||||
"consecutive": consecutive,
|
||||
"cooldown_seconds": self.cooldown_seconds,
|
||||
"notification_enabled": self.enabled,
|
||||
"notification_channels": self.channels,
|
||||
}
|
||||
|
||||
def reset(self, session_id: str) -> None:
|
||||
with self._lock:
|
||||
self._states.pop(session_id, None)
|
||||
|
||||
def _prune_sessions(self, now: float) -> None:
|
||||
expired = [
|
||||
session_id
|
||||
for session_id, state in self._states.items()
|
||||
if now - state.last_seen_at > self.session_ttl_seconds
|
||||
]
|
||||
for session_id in expired:
|
||||
del self._states[session_id]
|
||||
|
||||
def _send_alerts(
|
||||
self,
|
||||
image: Image.Image,
|
||||
triggered: list[str],
|
||||
detections: list[dict[str, Any]],
|
||||
) -> None:
|
||||
if self.wechat_webhook_url:
|
||||
try:
|
||||
self._send_wechat_alert(image, triggered, detections)
|
||||
except (HTTPError, URLError, TimeoutError, ValueError) as error:
|
||||
LOGGER.error("WeChat alert failed: %s", error)
|
||||
if self.feishu_webhook_url:
|
||||
try:
|
||||
self._send_feishu_alert(triggered, detections)
|
||||
except (HTTPError, URLError, TimeoutError, ValueError) as error:
|
||||
LOGGER.error("Feishu alert failed: %s", error)
|
||||
|
||||
def _send_wechat_alert(
|
||||
self,
|
||||
image: Image.Image,
|
||||
triggered: list[str],
|
||||
detections: list[dict[str, Any]],
|
||||
) -> None:
|
||||
if not self.wechat_webhook_url:
|
||||
return
|
||||
target_text, max_confidence = alert_summary(triggered, detections)
|
||||
message = (
|
||||
f"🔥 烟火检测告警\n"
|
||||
f"> 检测目标：{target_text}\n"
|
||||
f"> 最高置信度：{max_confidence:.1%}\n"
|
||||
f"> 请及时查看现场视频。"
|
||||
)
|
||||
post_json(
|
||||
self.wechat_webhook_url,
|
||||
{"msgtype": "markdown", "markdown": {"content": message}},
|
||||
)
|
||||
image_bytes = encode_jpeg(image)
|
||||
post_json(
|
||||
self.wechat_webhook_url,
|
||||
{
|
||||
"msgtype": "image",
|
||||
"image": {
|
||||
"base64": base64.b64encode(image_bytes).decode("ascii"),
|
||||
"md5": hashlib.md5(
|
||||
image_bytes,
|
||||
usedforsecurity=False,
|
||||
).hexdigest(),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
def _send_feishu_alert(
|
||||
self,
|
||||
triggered: list[str],
|
||||
detections: list[dict[str, Any]],
|
||||
) -> None:
|
||||
if not self.feishu_webhook_url:
|
||||
return
|
||||
target_text, max_confidence = alert_summary(triggered, detections)
|
||||
payload: dict[str, Any] = {
|
||||
"msg_type": "interactive",
|
||||
"card": {
|
||||
"config": {"wide_screen_mode": True},
|
||||
"header": {
|
||||
"template": "red",
|
||||
"title": {
|
||||
"tag": "plain_text",
|
||||
"content": "烟火检测告警",
|
||||
},
|
||||
},
|
||||
"elements": [
|
||||
{
|
||||
"tag": "markdown",
|
||||
"content": (
|
||||
f"**检测目标：** {target_text}\n"
|
||||
f"**最高置信度：** {max_confidence:.1%}\n"
|
||||
f"**告警时间：** "
|
||||
f"{time.strftime('%Y-%m-%d %H:%M:%S')}\n"
|
||||
f"请及时查看现场视频。"
|
||||
),
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
if self.feishu_secret:
|
||||
timestamp = str(int(time.time()))
|
||||
payload["timestamp"] = timestamp
|
||||
payload["sign"] = feishu_signature(
|
||||
timestamp,
|
||||
self.feishu_secret,
|
||||
)
|
||||
post_json(self.feishu_webhook_url, payload)
|
||||
|
||||
|
||||
def alert_summary(
|
||||
triggered: list[str],
|
||||
detections: list[dict[str, Any]],
|
||||
) -> tuple[str, float]:
|
||||
labels = {"fire": "火焰", "smoke": "烟雾"}
|
||||
target_text = "、".join(labels[name] for name in triggered)
|
||||
max_confidence = max(
|
||||
(
|
||||
float(detection.get("confidence", 0.0))
|
||||
for detection in detections
|
||||
if detection.get("class") in triggered
|
||||
),
|
||||
default=0.0,
|
||||
)
|
||||
return target_text, max_confidence
|
||||
|
||||
|
||||
def feishu_signature(timestamp: str, secret: str) -> str:
|
||||
string_to_sign = f"{timestamp}\n{secret}".encode("utf-8")
|
||||
digest = hmac.new(string_to_sign, digestmod=hashlib.sha256).digest()
|
||||
return base64.b64encode(digest).decode("ascii")
|
||||
|
||||
|
||||
def post_json(url: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
request = Request(
|
||||
url,
|
||||
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urlopen(request, timeout=10) as response:
|
||||
result = json.loads(response.read().decode("utf-8"))
|
||||
error_code = result.get(
|
||||
"errcode",
|
||||
result.get("code", result.get("StatusCode", 0)),
|
||||
)
|
||||
if error_code != 0:
|
||||
raise ValueError(
|
||||
result.get("errmsg")
|
||||
or result.get("msg")
|
||||
or result.get("StatusMessage")
|
||||
or "Unknown robot webhook error"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def annotate_image(
|
||||
image: Image.Image,
|
||||
detections: list[dict[str, Any]],
|
||||
) -> Image.Image:
|
||||
annotated = image.copy()
|
||||
draw = ImageDraw.Draw(annotated)
|
||||
font = ImageFont.load_default()
|
||||
colors = {"fire": "#ff3b30", "smoke": "#00a89b"}
|
||||
for detection in detections:
|
||||
box = detection.get("box")
|
||||
if not box or len(box) != 4:
|
||||
continue
|
||||
class_name = str(detection.get("class", "target"))
|
||||
confidence = float(detection.get("confidence", 0.0))
|
||||
color = colors.get(class_name, "#ffd166")
|
||||
coordinates = tuple(int(round(value)) for value in box)
|
||||
draw.rectangle(coordinates, outline=color, width=4)
|
||||
draw.text(
|
||||
(coordinates[0] + 4, max(0, coordinates[1] - 16)),
|
||||
f"{class_name} {confidence:.0%}",
|
||||
fill=color,
|
||||
font=font,
|
||||
)
|
||||
return annotated
|
||||
|
||||
|
||||
def encode_jpeg(image: Image.Image, max_bytes: int = 1_900_000) -> bytes:
|
||||
working = image.convert("RGB")
|
||||
quality = 88
|
||||
while True:
|
||||
output = BytesIO()
|
||||
working.save(output, format="JPEG", quality=quality, optimize=True)
|
||||
payload = output.getvalue()
|
||||
if len(payload) <= max_bytes:
|
||||
return payload
|
||||
if quality > 55:
|
||||
quality -= 10
|
||||
continue
|
||||
width, height = working.size
|
||||
working = working.resize(
|
||||
(max(1, int(width * 0.8)), max(1, int(height * 0.8)))
|
||||
)
|
||||
quality = 75
|
||||
Reference in New Issue
Block a user