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 pydantic import BaseModel from ultralytics import YOLO from .alerting import AlertManager from .events import EVENT_STATUSES, EventStore 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")), ) EVENT_STORE = EventStore() class EventStatusUpdate(BaseModel): status: str 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, } ) if result["alert"]["triggered"]: result["event"] = EVENT_STORE.create( session_id=session_id or "single-image", classes=result["alert"]["classes"], detections=result["detections"], notification_channels=result["alert"].get( "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("/api/events") def list_events( limit: int = 100, status: str | None = None, ) -> dict[str, Any]: if status and status not in EVENT_STATUSES: raise HTTPException(status_code=400, detail="Invalid event status") return { "events": EVENT_STORE.list(limit=limit, status=status), "summary": EVENT_STORE.summary(), } @app.patch("/api/events/{event_id}") def update_event( event_id: str, update: EventStatusUpdate, ) -> dict[str, Any]: try: event = EVENT_STORE.update(event_id, update.status) except ValueError as error: raise HTTPException(status_code=400, detail=str(error)) from error if event is None: raise HTTPException(status_code=404, detail="Event not found") return event @app.get("/api/dashboard") def dashboard() -> dict[str, Any]: return { "summary": EVENT_STORE.summary(), "recent_events": EVENT_STORE.list(limit=6), "system": health(), "detection": { "confidence": CONFIDENCE, "iou": IOU, "image_size": IMAGE_SIZE, "confirm_frames": ALERT_MANAGER.confirm_frames, "cooldown_seconds": ALERT_MANAGER.cooldown_seconds, }, } @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)