324 lines
10 KiB
Python
324 lines
10 KiB
Python
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
|