135 lines
4.2 KiB
Python
135 lines
4.2 KiB
Python
from fastapi import FastAPI, UploadFile, File, WebSocket, WebSocketDisconnect
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from contextlib import asynccontextmanager
|
||
import uvicorn
|
||
import os
|
||
import signal
|
||
import sys
|
||
import logging
|
||
|
||
from api import detection, models
|
||
from services.model_service import ModelService
|
||
from services.camera_service import CameraService
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
model_service = ModelService()
|
||
camera_service = None
|
||
|
||
|
||
def setup_signal_handlers():
|
||
"""设置信号处理器，确保进程异常退出时能清理资源"""
|
||
def signal_handler(signum, frame):
|
||
sig_name = signal.Signals(signum).name
|
||
logger.info(f"收到信号 {sig_name}，正在清理资源...")
|
||
|
||
# 强制释放摄像头资源
|
||
import subprocess
|
||
try:
|
||
# 查找并终止占用摄像头的Python进程（除了当前进程）
|
||
current_pid = os.getpid()
|
||
result = subprocess.run(
|
||
['lsof', '+D', '/dev', '-a', '-c', 'python'],
|
||
capture_output=True,
|
||
text=True
|
||
)
|
||
if result.returncode == 0:
|
||
for line in result.stdout.split('\n'):
|
||
if '/dev/video' in line:
|
||
parts = line.split()
|
||
if len(parts) >= 2:
|
||
try:
|
||
pid = int(parts[1])
|
||
if pid != current_pid:
|
||
logger.info(f"终止占用摄像头的进程: {pid}")
|
||
os.kill(pid, signal.SIGTERM)
|
||
except (ValueError, ProcessLookupError):
|
||
pass
|
||
except Exception as e:
|
||
logger.error(f"清理摄像头资源失败: {e}")
|
||
|
||
sys.exit(0)
|
||
|
||
# 注册信号处理器
|
||
signal.signal(signal.SIGTERM, signal_handler)
|
||
signal.signal(signal.SIGINT, signal_handler)
|
||
if hasattr(signal, 'SIGQUIT'):
|
||
signal.signal(signal.SIGQUIT, signal_handler)
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
global camera_service
|
||
|
||
camera_service = CameraService(model_service)
|
||
yield
|
||
|
||
# 关闭时清理资源
|
||
logger.info("正在关闭服务，清理资源...")
|
||
if camera_service:
|
||
await camera_service.stop()
|
||
|
||
app = FastAPI(
|
||
title="视频模型检测平台",
|
||
description="基于YOLO的实时视频检测平台",
|
||
version="1.0.0",
|
||
lifespan=lifespan
|
||
)
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
app.mount("/static", StaticFiles(directory="static"), name="static")
|
||
|
||
docker_output_dir = os.path.join(
|
||
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
|
||
"external", "video-recognition-system", "PaddlePaddle", "PaddleDetection", "output"
|
||
)
|
||
os.makedirs(docker_output_dir, exist_ok=True)
|
||
app.mount("/docker-output", StaticFiles(directory=docker_output_dir), name="docker-output")
|
||
|
||
app.include_router(detection.router, prefix="/api")
|
||
app.include_router(models.router, prefix="/api")
|
||
|
||
@app.get("/")
|
||
async def root():
|
||
return {"message": "视频模型检测平台 API", "version": "1.0.0"}
|
||
|
||
@app.get("/api/health")
|
||
async def health():
|
||
return {"status": "healthy"}
|
||
|
||
@app.websocket("/ws/camera")
|
||
async def camera_websocket_endpoint(websocket: WebSocket):
|
||
await camera_service.handle_connection(websocket)
|
||
|
||
if __name__ == "__main__":
|
||
os.makedirs("static/uploads", exist_ok=True)
|
||
os.makedirs("static/results", exist_ok=True)
|
||
os.makedirs("static/temp", exist_ok=True)
|
||
|
||
# 设置信号处理器
|
||
setup_signal_handlers()
|
||
|
||
# 检测是否处于uvicorn重载模式的子进程中
|
||
is_reload_worker = os.environ.get('UVICORN_RELOAD') == 'true'
|
||
|
||
if is_reload_worker:
|
||
logger.info("检测到uvicorn重载子进程，跳过摄像头预清理")
|
||
|
||
uvicorn.run(
|
||
"main:app",
|
||
host="0.0.0.0",
|
||
port=8000,
|
||
reload=True,
|
||
reload_dirs=["./"],
|
||
reload_includes=["*.py"]
|
||
)
|