Initial commit: Video detection platform with YOLO models
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
This commit is contained in:
67
apps/server/api/detection.py
Normal file
67
apps/server/api/detection.py
Normal file
@@ -0,0 +1,67 @@
|
||||
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={}
|
||||
)
|
||||
Reference in New Issue
Block a user