8fb58c75fe
Features: - Fire detection (YOLOv10) - Helmet detection (YOLOv8) - Crowd detection (YOLOv8) - Smoking detection (YOLOv8) - Loitering detection (YOLOv8) Tech Stack: - Frontend: Vue 3 + Vite + Element Plus - Backend: FastAPI + WebSocket - Monorepo: pnpm workspace + Turbo - Docker support included
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
import base64
|
|
import logging
|
|
from fastapi import APIRouter, UploadFile, File, Form, Query
|
|
from models.schemas import ImageDetectionResult
|
|
|
|
router = APIRouter()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
@router.post("/detect/image", response_model=ImageDetectionResult)
|
|
async def detect_image(
|
|
file: UploadFile = File(...),
|
|
model_id: str = Query("fire_detection"),
|
|
confidence: float = Query(0.5),
|
|
iou: float = Query(0.45)
|
|
):
|
|
from main import model_service
|
|
from services.detection_service import DetectionService
|
|
|
|
detection_service = DetectionService(model_service)
|
|
|
|
try:
|
|
contents = await file.read()
|
|
nparr = np.frombuffer(contents, np.uint8)
|
|
frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
|
|
if frame is None:
|
|
return ImageDetectionResult(
|
|
success=False,
|
|
message="无法读取图片",
|
|
data={}
|
|
)
|
|
|
|
result = await detection_service.detect_image(frame, model_id, confidence, iou)
|
|
|
|
if result['success']:
|
|
annotated_frame = detection_service.draw_detections(frame, result['detections'])
|
|
|
|
import uuid
|
|
result_filename = f"result_{uuid.uuid4().hex[:8]}.jpg"
|
|
result_path = f"static/results/{result_filename}"
|
|
cv2.imwrite(result_path, annotated_frame)
|
|
|
|
return ImageDetectionResult(
|
|
success=True,
|
|
message="检测完成",
|
|
data={
|
|
"detections": result['detections'],
|
|
"image_url": f"/static/results/{result_filename}",
|
|
"stats": result['stats']
|
|
}
|
|
)
|
|
else:
|
|
return ImageDetectionResult(
|
|
success=False,
|
|
message=result['message'],
|
|
data={}
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"图片检测失败: {e}")
|
|
return ImageDetectionResult(
|
|
success=False,
|
|
message=f"检测失败: {str(e)}",
|
|
data={}
|
|
)
|