215 lines
6.8 KiB
Python
215 lines
6.8 KiB
Python
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)
|