#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 入户门检测 FastAPI 服务 提供 HTTP API 接口用于入户门位置检测： - POST /detect - 检测入户门位置 - GET /status/{task_id} - 查询任务状态 - GET /result/{task_id} - 获取检测结果 - GET /health - 健康检查 """ import os import sys import json import shutil import uuid import asyncio from pathlib import Path from typing import Dict, List, Optional, Any from datetime import datetime from concurrent.futures import ThreadPoolExecutor import tempfile from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File, Form from fastapi.responses import JSONResponse, FileResponse from pydantic import BaseModel, Field import uvicorn # 导入入户门检测器 from export_entrance_position import EntranceDoorDetector # ============================================================================ # 数据模型 # ============================================================================ class DetectRequest(BaseModel): """检测请求体""" scene_folder: str = Field(..., description="场景文件夹路径") model_path: Optional[str] = Field("yoloe-26x-seg.pt", description="YOLOE 模型路径") conf: Optional[float] = Field(0.35, description="检测置信度阈值", ge=0, le=1) iou: Optional[float] = Field(0.45, description="NMS IoU 阈值", ge=0, le=1) voxel_size: Optional[float] = Field(0.03, description="点云体素尺寸") imgsz: Optional[List[int]] = Field([1024, 2048], description="YOLOE 输入图像尺寸 [高，宽]") vis_ply: Optional[bool] = Field(False, description="是否生成可视化 PLY") class DetectResponse(BaseModel): """检测响应体""" task_id: str status: str message: str result_path: Optional[str] = None vis_ply_path: Optional[str] = None class StatusResponse(BaseModel): """任务状态响应""" task_id: str status: str # pending, processing, completed, failed progress: Optional[float] = None # 0-1 message: Optional[str] = None result: Optional[Dict[str, Any]] = None error: Optional[str] = None created_at: str completed_at: Optional[str] = None class HealthResponse(BaseModel): """健康检查响应""" status: str model_loaded: bool gpu_available: bool timestamp: str # ============================================================================ # 任务管理 # ============================================================================ class TaskManager: """任务管理器""" def __init__(self): self.tasks: Dict[str, Dict[str, Any]] = {} self.results_dir = Path("server_results") self.results_dir.mkdir(exist_ok=True) def create_task(self, task_id: str, scene_folder: str) -> Dict[str, Any]: """创建新任务""" task = { "task_id": task_id, "status": "pending", "progress": 0.0, "message": "任务已创建，等待处理", "result": None, "error": None, "created_at": datetime.now().isoformat(), "completed_at": None, "scene_folder": scene_folder, "result_path": None, "vis_ply_path": None } self.tasks[task_id] = task return task def update_task(self, task_id: str, **kwargs): """更新任务状态""" if task_id not in self.tasks: raise ValueError(f"任务 {task_id} 不存在") for key, value in kwargs.items(): self.tasks[task_id][key] = value def get_task(self, task_id: str) -> Optional[Dict[str, Any]]: """获取任务信息""" return self.tasks.get(task_id) def get_task_status(self, task_id: str) -> StatusResponse: """获取任务状态""" task = self.tasks.get(task_id) if not task: raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在") return StatusResponse( task_id=task["task_id"], status=task["status"], progress=task["progress"], message=task["message"], result=task["result"], error=task["error"], created_at=task["created_at"], completed_at=task["completed_at"] ) # ============================================================================ # FastAPI 应用 # ============================================================================ app = FastAPI( title="入户门检测服务", description="基于 YOLOE 和 RGB-D 点云的入户门位置检测 API", version="1.0.0" ) # 任务管理器 task_manager = TaskManager() # 线程池 executor = ThreadPoolExecutor(max_workers=2) # 模型加载状态 model_loaded = False detector_instance = None @app.on_event("startup") async def startup_event(): """服务启动时预加载模型""" global model_loaded, detector_instance model_path = "yoloe-26x-seg.pt" if Path(model_path).exists(): try: # 预加载模型（创建一个临时检测器实例） print(f"预加载 YOLOE 模型：{model_path}") # 注意：这里不实际初始化，因为需要场景文件夹 model_loaded = True print("模型加载成功") except Exception as e: print(f"模型加载失败：{e}") model_loaded = False else: print(f"模型文件不存在：{model_path}") model_loaded = False @app.get("/health", response_model=HealthResponse) async def health_check(): """健康检查""" import torch return HealthResponse( status="healthy", model_loaded=model_loaded, gpu_available=torch.cuda.is_available(), timestamp=datetime.now().isoformat() ) @app.post("/detect", response_model=DetectResponse) async def detect_entrance_door( request: DetectRequest, background_tasks: BackgroundTasks ): """ 检测入户门位置 提交检测任务后，使用 task_id 查询状态和结果 """ # 验证场景文件夹 scene_path = Path(request.scene_folder) if not scene_path.exists(): raise HTTPException(status_code=400, detail=f"场景文件夹不存在：{request.scene_folder}") # 创建任务 task_id = str(uuid.uuid4()) task_manager.create_task(task_id, request.scene_folder) # 异步执行检测 background_tasks.add_task( run_detection, task_id=task_id, scene_folder=request.scene_folder, model_path=request.model_path, conf=request.conf, iou=request.iou, voxel_size=request.voxel_size, imgsz=request.imgsz, vis_ply=request.vis_ply ) return DetectResponse( task_id=task_id, status="pending", message="任务已提交，正在处理中", result_path=None, vis_ply_path=None ) @app.post("/detect/upload", response_model=DetectResponse) async def detect_entrance_door_upload( scene_folder: str = Form(...), model_path: Optional[str] = Form("yoloe-26x-seg.pt"), conf: Optional[float] = Form(0.35), iou: Optional[float] = Form(0.45), voxel_size: Optional[float] = Form(0.03), imgsz_height: Optional[int] = Form(1024), imgsz_width: Optional[int] = Form(2048), vis_ply: Optional[bool] = Form(False), files: List[UploadFile] = File(...), background_tasks: BackgroundTasks = None ): """ 上传场景文件并检测入户门位置 需要上传： - vision.txt 文件 - pano_img/ 文件夹中的 JPG 图像 - depth_img/ 文件夹中的 PNG 深度图 """ # 创建临时目录 temp_dir = tempfile.mkdtemp(prefix="entrance_detect_") temp_path = Path(temp_dir) try: # 保存上传的文件 pano_dir = temp_path / "pano_img" depth_dir = temp_path / "depth_img" pano_dir.mkdir() depth_dir.mkdir() for file in files: if file.filename.endswith(".jpg"): file_path = pano_dir / file.filename elif file.filename.endswith(".png"): file_path = depth_dir / file.filename elif file.filename == "vision.txt": file_path = temp_path / "vision.txt" else: continue with open(file_path, "wb") as f: content = await file.read() f.write(content) # 创建任务 task_id = str(uuid.uuid4()) task_manager.create_task(task_id, str(temp_path)) # 异步执行检测 background_tasks.add_task( run_detection, task_id=task_id, scene_folder=str(temp_path), model_path=model_path, conf=conf, iou=iou, voxel_size=voxel_size, imgsz=[imgsz_height, imgsz_width], vis_ply=vis_ply, cleanup_temp=True ) return DetectResponse( task_id=task_id, status="pending", message="文件上传成功，任务已提交", result_path=None, vis_ply_path=None ) except Exception as e: # 清理临时目录 shutil.rmtree(temp_dir, ignore_errors=True) raise HTTPException(status_code=500, detail=f"文件上传失败：{str(e)}") @app.get("/status/{task_id}", response_model=StatusResponse) async def get_task_status(task_id: str): """查询任务状态""" return task_manager.get_task_status(task_id) @app.get("/result/{task_id}") async def get_task_result(task_id: str): """获取检测结果""" task = task_manager.get_task(task_id) if not task: raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在") if task["status"] != "completed": raise HTTPException( status_code=400, detail=f"任务尚未完成，当前状态：{task['status']}" ) # 读取结果文件 if task["result_path"] and Path(task["result_path"]).exists(): with open(task["result_path"], "r", encoding="utf-8") as f: result = json.load(f) return JSONResponse(content=result) else: raise HTTPException(status_code=404, detail="结果文件不存在") @app.get("/result/{task_id}/ply") async def get_task_ply(task_id: str): """获取可视化 PLY 文件""" task = task_manager.get_task(task_id) if not task: raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在") if task["status"] != "completed": raise HTTPException( status_code=400, detail=f"任务尚未完成，当前状态：{task['status']}" ) if task["vis_ply_path"] and Path(task["vis_ply_path"]).exists(): return FileResponse( task["vis_ply_path"], media_type="application/octet-stream", filename="vis.ply" ) else: raise HTTPException(status_code=404, detail="PLY 文件不存在") @app.delete("/task/{task_id}") async def delete_task(task_id: str): """删除任务""" if task_id not in task_manager.tasks: raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在") task = task_manager.tasks[task_id] # 清理结果文件 if task["result_path"] and Path(task["result_path"]).exists(): os.remove(task["result_path"]) if task["vis_ply_path"] and Path(task["vis_ply_path"]).exists(): os.remove(task["vis_ply_path"]) # 清理场景文件夹（如果是临时创建的） scene_path = Path(task.get("scene_folder", "")) if scene_path.exists() and str(scene_path).startswith("/tmp"): shutil.rmtree(scene_path, ignore_errors=True) del task_manager.tasks[task_id] return {"message": f"任务 {task_id} 已删除"} # ============================================================================ # 后台任务 # ============================================================================ def run_detection( task_id: str, scene_folder: str, model_path: str, conf: float, iou: float, voxel_size: float, imgsz: List[int], vis_ply: bool, cleanup_temp: bool = False ): """运行入户门检测（在后台线程中执行）""" try: # 更新状态为处理中 task_manager.update_task( task_id, status="processing", progress=0.1, message="正在初始化检测器..." ) # 创建检测器 detector = EntranceDoorDetector( scene_folder=scene_folder, model_path=model_path, conf=conf, iou=iou, voxel_size=voxel_size, imgsz=(imgsz[0], imgsz[1]) ) task_manager.update_task( task_id, progress=0.3, message="正在检测门并识别入户门..." ) # 执行检测 success = detector.detect_and_identify() task_manager.update_task( task_id, progress=0.7, message="正在导出结果..." ) if success: # 导出 JSON 结果 output_folder = Path(scene_folder) / "output" output_folder.mkdir(exist_ok=True) result_path = output_folder / "entrance_position.json" detector.export_json(str(result_path)) vis_ply_path = None if vis_ply: vis_ply_path = output_folder / "vis.ply" detector.export_vis_ply(detector.combined_pc, str(vis_ply_path)) # 读取结果 with open(result_path, "r", encoding="utf-8") as f: result = json.load(f) # 更新任务为完成 task_manager.update_task( task_id, status="completed", progress=1.0, message="检测完成", result=result, result_path=str(result_path), vis_ply_path=str(vis_ply_path) if vis_ply_path else None, completed_at=datetime.now().isoformat() ) else: task_manager.update_task( task_id, status="failed", progress=0.0, message="检测失败：无法确定入户门位置", error="无法确定入户门位置", completed_at=datetime.now().isoformat() ) # 清理临时目录 if cleanup_temp: shutil.rmtree(scene_folder, ignore_errors=True) except Exception as e: import traceback error_msg = f"{str(e)}\n{traceback.format_exc()}" task_manager.update_task( task_id, status="failed", progress=0.0, message=f"检测失败：{str(e)}", error=error_msg, completed_at=datetime.now().isoformat() ) # ============================================================================ # 主程序 # ============================================================================ if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="入户门检测 FastAPI 服务") parser.add_argument("--host", type=str, default="0.0.0.0", help="监听地址") parser.add_argument("--port", type=int, default=8000, help="监听端口") parser.add_argument("--reload", action="store_true", help="启用热重载（开发模式）") parser.add_argument("--workers", type=int, default=1, help="工作进程数") args = parser.parse_args() print(f"启动入户门检测服务...") print(f" 监听地址：http://{args.host}:{args.port}") print(f" API 文档：http://{args.host}:{args.port}/docs") print(f" 健康检查：http://{args.host}:{args.port}/health") uvicorn.run( "server:app", host=args.host, port=args.port, reload=args.reload, workers=args.workers )