feat: add video detection and robot alerts
This commit is contained in:
1
backend/__init__.py
Normal file
1
backend/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""HTTP inference backend for the smoke and fire detector."""
|
||||
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
|
||||
154
backend/main.py
Normal file
154
backend/main.py
Normal file
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import os
|
||||
import time
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, File, HTTPException, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
from dotenv import load_dotenv
|
||||
from ultralytics import YOLO
|
||||
|
||||
from .alerting import AlertManager
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
FRONTEND_DIR = PROJECT_ROOT / "frontend"
|
||||
load_dotenv(PROJECT_ROOT / ".env")
|
||||
DEFAULT_WEIGHTS = PROJECT_ROOT / "runs" / "detect" / "smoke_fire_yolo11s_v1-4" / "weights" / "best.pt"
|
||||
WEIGHTS_PATH = Path(os.getenv("YOLO_WEIGHTS", str(DEFAULT_WEIGHTS))).expanduser().resolve()
|
||||
DEVICE = os.getenv("YOLO_DEVICE") or None
|
||||
IMAGE_SIZE = int(os.getenv("YOLO_IMGSZ", "768"))
|
||||
CONFIDENCE = float(os.getenv("YOLO_CONF", "0.40"))
|
||||
IOU = float(os.getenv("YOLO_IOU", "0.45"))
|
||||
MAX_UPLOAD_BYTES = 15 * 1024 * 1024
|
||||
MAX_IMAGE_PIXELS = int(os.getenv("YOLO_MAX_IMAGE_PIXELS", "25000000"))
|
||||
CLASS_NAMES = {0: "smoke", 1: "fire"}
|
||||
ALERT_MANAGER = AlertManager(
|
||||
confirm_frames=int(os.getenv("ALERT_CONFIRM_FRAMES", "3")),
|
||||
cooldown_seconds=float(os.getenv("ALERT_COOLDOWN_SECONDS", "60")),
|
||||
)
|
||||
|
||||
app = FastAPI(title="Smoke Fire Detector API", version="0.1.0")
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_methods=["GET", "POST", "DELETE"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_model() -> YOLO:
|
||||
if not WEIGHTS_PATH.is_file():
|
||||
raise FileNotFoundError(f"YOLO weights not found: {WEIGHTS_PATH}")
|
||||
return YOLO(str(WEIGHTS_PATH))
|
||||
|
||||
|
||||
def validate_image(payload: bytes) -> Image.Image:
|
||||
if not payload:
|
||||
raise HTTPException(status_code=400, detail="Uploaded file is empty")
|
||||
if len(payload) > MAX_UPLOAD_BYTES:
|
||||
raise HTTPException(status_code=413, detail="Uploaded file exceeds 15 MB")
|
||||
try:
|
||||
image = Image.open(io.BytesIO(payload))
|
||||
if image.width * image.height > MAX_IMAGE_PIXELS:
|
||||
raise HTTPException(
|
||||
status_code=413,
|
||||
detail="Image dimensions are too large",
|
||||
)
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
except (OSError, UnidentifiedImageError) as error:
|
||||
raise HTTPException(status_code=415, detail="Only valid image files are supported") from error
|
||||
|
||||
|
||||
def predict_image(image: Image.Image) -> dict[str, Any]:
|
||||
started_at = time.perf_counter()
|
||||
try:
|
||||
result = get_model().predict(
|
||||
source=image,
|
||||
conf=CONFIDENCE,
|
||||
iou=IOU,
|
||||
imgsz=IMAGE_SIZE,
|
||||
device=DEVICE,
|
||||
verbose=False,
|
||||
)[0]
|
||||
except FileNotFoundError as error:
|
||||
raise HTTPException(status_code=503, detail=str(error)) from error
|
||||
except Exception as error:
|
||||
raise HTTPException(status_code=500, detail=f"Inference failed: {error}") from error
|
||||
detections = []
|
||||
for box in result.boxes:
|
||||
class_id = int(box.cls.item())
|
||||
detections.append({
|
||||
"class": CLASS_NAMES.get(class_id, str(class_id)),
|
||||
"class_id": class_id,
|
||||
"confidence": round(float(box.conf.item()), 6),
|
||||
"box": [round(float(value), 2) for value in box.xyxy[0].tolist()],
|
||||
})
|
||||
return {
|
||||
"detections": detections,
|
||||
"image": {"width": image.width, "height": image.height},
|
||||
"inference_ms": round((time.perf_counter() - started_at) * 1000, 1),
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/health")
|
||||
def health() -> dict[str, Any]:
|
||||
return {
|
||||
"status": "ok",
|
||||
"weights": str(WEIGHTS_PATH),
|
||||
"weights_available": WEIGHTS_PATH.is_file(),
|
||||
"wechat_alerts_enabled": ALERT_MANAGER.channels["wechat"],
|
||||
"feishu_alerts_enabled": ALERT_MANAGER.channels["feishu"],
|
||||
"alert_channels": ALERT_MANAGER.channels,
|
||||
}
|
||||
|
||||
|
||||
@app.post("/api/detect")
|
||||
async def detect(
|
||||
file: UploadFile = File(...),
|
||||
session_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
payload = await file.read(MAX_UPLOAD_BYTES + 1)
|
||||
image = validate_image(payload)
|
||||
result = predict_image(image)
|
||||
result["alert"] = (
|
||||
ALERT_MANAGER.evaluate(
|
||||
session_id,
|
||||
result["detections"],
|
||||
image,
|
||||
)
|
||||
if session_id
|
||||
else {
|
||||
"triggered": False,
|
||||
"classes": [],
|
||||
"notification_enabled": ALERT_MANAGER.enabled,
|
||||
"notification_channels": ALERT_MANAGER.channels,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@app.delete("/api/sessions/{session_id}")
|
||||
def reset_detection_session(session_id: str) -> dict[str, str]:
|
||||
ALERT_MANAGER.reset(session_id)
|
||||
return {"status": "reset"}
|
||||
|
||||
|
||||
@app.get("/")
|
||||
def frontend() -> FileResponse:
|
||||
return FileResponse(FRONTEND_DIR / "index.html")
|
||||
|
||||
|
||||
@app.get("/{asset_path:path}")
|
||||
def frontend_asset(asset_path: str) -> FileResponse:
|
||||
requested = (FRONTEND_DIR / asset_path).resolve()
|
||||
if FRONTEND_DIR not in requested.parents or not requested.is_file():
|
||||
raise HTTPException(status_code=404, detail="Asset not found")
|
||||
return FileResponse(requested)
|
||||
Reference in New Issue
Block a user