hongchen/one hundred_class/api_server_with_result.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
房间分割 & 家具检测 FastAPI 服务（返回合并后的家具检测数据）

启动：
    python api_server_with_result.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 json
import os
import sys
import traceback
from contextlib import asynccontextmanager
from pathlib import Path

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="输出目录路径映射"
    )
    furniture_data: list = Field(
        default_factory=list,
        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,
)


# ============================================================================
# 辅助函数
# ============================================================================

def merge_furniture_jsons(detect_dir: str) -> list:
    """读取 detect_dir 下所有 .json 文件，按文件名排序后合并为一个 list 返回。

    每个元素保留原始 JSON 的完整结构（shapes、imagePath、imageWidth、
    imageHeight、version 以及 Step3 注入的 belong_room 等字段）。
    """
    result = []
    json_files = sorted(
        Path(detect_dir).glob("*.json"),
        key=lambda p: _natural_sort_key(p.stem),
    )
    for jf in json_files:
        with open(jf, "r", encoding="utf-8") as f:
            result.append(json.load(f))
    return result


def _natural_sort_key(name: str):
    """将字符串中的数字部分转为 int 元组，用于自然排序（0, 1, 2, ..., 10, 11）。"""
    import re
    return [
        int(part) if part.isdigit() else part.lower()
        for part in re.split(r"(\d+)", name)
    ]


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

@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")
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)}")

    # --- 合并家具检测 JSON 并返回 ---
    print(f"\n[API] >>> 合并家具检测数据：{out_detect_dir}")
    furniture_data = merge_furniture_jsons(out_detect_dir)

    print(f"\n[API] >>> 处理完成：{scene_path}，共 {len(furniture_data)} 张图片的检测结果")

    return {
        "status": "success",
        "message": "处理完成",
        "scene_path": scene_path,
        "output_dirs": {
            "room_output": out_room_dir,
            "furniture_output": out_detect_dir,
        },
        "furniture_data": furniture_data,
    }


@app.get("/docs")
def swagger_ui():
    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_with_result:app",
        host=args.host,
        port=args.port,
        workers=args.workers,
        log_level="info",
    )