hongchen/one hundred_class/api_server.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
房间分割 & 家具检测 FastAPI 服务

启动：
    python api_server.py --host 0.0.0.0 --port 8000

调用示例：
    curl -X POST http://localhost:8000/api/process \
         -H "Content-Type: application/json" \
         -d '{"scene_path": "/absolute/path/to/SG-xxxxx"}'
"""
import asyncio
import json
import os
import sys
import traceback
from contextlib import asynccontextmanager
from typing import Optional

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field


# ============================================================================
# 确保 pipeline 模块可用
# ============================================================================
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import pipeline  # noqa: E402


# ============================================================================
# 请求 / 响应模型
# ============================================================================

class ProcessRequest(BaseModel):
    scene_path: str = Field(
        ...,
        description="场景文件夹的绝对路径（包含 images/、depthmap/、vision.txt 等）",
        examples=["/data/scenes/SG-n6nV8B2oW95"]
    )


class ProcessResponse(BaseModel):
    status: str = Field(..., description="success / error")
    message: str
    scene_path: str
    output_dirs: dict = Field(
        default_factory=dict,
        description="输出目录路径映射"
    )


class HealthResponse(BaseModel):
    status: str
    model_dir: str


# ============================================================================
# 启动 / 关闭
# ============================================================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    model_dir = os.path.dirname(os.path.abspath(__file__))
    print(f"[API] 模型目录：{model_dir}")

    for name, filename in [
        ("BiRefNet 权重", "initial_mask.onnx"),
        ("房间分类模型", "room_cls.pt"),
        ("家具检测模型", "100_cls_seg.pt"),
    ]:
        path = os.path.join(model_dir, filename)
        if not os.path.exists(path):
            print(f"[API WARNING] 模型文件不存在：{name} -> {path}")
        else:
            print(f"[API OK] 模型就绪：{filename}")

    yield
    print("[API] 服务已关闭")


app = FastAPI(
    title="Room & Furniture Pipeline API",
    description="房间分割 → 家具检测 → 家具映射到房间",
    version="1.0.0",
    lifespan=lifespan,
)


# ============================================================================
# 路由
# ============================================================================

@app.get("/health", response_model=HealthResponse)
def health_check():
    model_dir = os.path.dirname(os.path.abspath(__file__))
    return HealthResponse(status="ok", model_dir=model_dir)


@app.post("/api/process", response_model=ProcessResponse)
def process_scene(req: ProcessRequest):
    scene_path = req.scene_path.strip()

    # --- 验证 ---
    if not os.path.isabs(scene_path):
        raise HTTPException(status_code=400, detail="scene_path 必须是绝对路径")
    if not os.path.isdir(scene_path):
        raise HTTPException(status_code=400, detail=f"文件夹不存在：{scene_path}")

    for sub, name in [("images", "图片文件夹"), ("depthmap", "深度图文件夹")]:
        p = os.path.join(scene_path, sub)
        if not os.path.exists(p):
            raise HTTPException(status_code=400, detail=f"{name}不存在：{p}")

    vision_file = os.path.join(scene_path, "vision.txt")
    if not os.path.exists(vision_file):
        raise HTTPException(status_code=400, detail=f"位姿文件不存在：{vision_file}")

    model_dir = os.path.dirname(os.path.abspath(__file__))
    weights_path = os.path.join(model_dir, "initial_mask.onnx")
    room_cls_path = os.path.join(model_dir, "room_cls.pt")
    detect_model_path = os.path.join(model_dir, "100_cls_seg.pt")

    for model_p, model_name in [(weights_path, "BiRefNet 权重"), (room_cls_path, "房间分类模型"), (detect_model_path, "家具检测模型")]:
        if not os.path.exists(model_p):
            raise HTTPException(status_code=500, detail=f"模型文件不存在：{model_name} -> {model_p}")

    # --- 输出目录 ---
    out_room_dir = os.path.join(scene_path, "out_room")
    out_detect_dir = os.path.join(scene_path, "out_detect_furniture")

    # --- 执行 ---
    try:
        # Step 1
        print(f"\n[API] >>> Step 1: 房间分割 | {scene_path}")
        pipeline.step1_extract_room(scene_path, out_room_dir, weights_path, room_cls_path)

        # Step 2
        print(f"\n[API] >>> Step 2: 家具检测")
        pipeline.step2_detect_furniture(scene_path, out_detect_dir, detect_model_path)

        # Step 3
        print(f"\n[API] >>> Step 3: 家具映射到房间")
        pipeline.step3_map_furniture_to_room(
            scene_folder=scene_path,
            images_folder=os.path.join(scene_path, "images"),
            depth_folder=os.path.join(scene_path, "depthmap"),
            vision_file=vision_file,
            furniture_folder=out_detect_dir,
            room_folder=out_room_dir,
        )
    except Exception as e:
        print(f"\n[API ERROR] 处理失败：{scene_path}")
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"处理失败：{str(e)}")

    print(f"\n[API] >>> 处理完成：{scene_path}")

    return ProcessResponse(
        status="success",
        message="处理完成",
        scene_path=scene_path,
        output_dirs={
            "room_output": out_room_dir,
            "furniture_output": out_detect_dir,
        }
    )


@app.get("/docs")
def swagger_ui():
    # FastAPI 自动提供 /docs，这里留个兜底
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/docs")


# ============================================================================
# 启动入口
# ============================================================================

def parse_args():
    import argparse
    parser = argparse.ArgumentParser(description="启动 Pipeline FastAPI 服务")
    parser.add_argument("--host", default="0.0.0.0", help="监听地址 (默认 0.0.0.0)")
    parser.add_argument("--port", type=int, default=8000, help="监听端口 (默认 8000)")
    parser.add_argument("--workers", type=int, default=1, help="工作进程数 (默认 1，长任务建议用 1)")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    uvicorn.run(
        "api_server:app",
        host=args.host,
        port=args.port,
        workers=args.workers,
        log_level="info",
    )