feat: add video detection and robot alerts

backend/main.py

@@ -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)