Tool_Agent/tools/local/image_stitcher/main.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import List, Literal
from stitch_core import stitch_images

app = FastAPI(title="Image Stitcher API")


class StitchRequest(BaseModel):
    images: List[str] = Field(..., min_length=2, description="Base64 编码的图片列表（至少 2 张）")
    direction: Literal["horizontal", "vertical", "grid"] = "horizontal"
    columns: int = Field(2, ge=1)
    spacing: int = Field(0, ge=0)
    background_color: str = "#FFFFFF"
    resize_mode: Literal["none", "fit_width", "fit_height"] = "none"


class StitchResponse(BaseModel):
    image: str
    width: int
    height: int


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/stitch", response_model=StitchResponse)
def stitch(req: StitchRequest):
    try:
        result = stitch_images(
            images=req.images,
            direction=req.direction,
            columns=req.columns,
            spacing=req.spacing,
            background_color=req.background_color,
            resize_mode=req.resize_mode
        )
        return result
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"拼接失败: {str(e)}")


if __name__ == "__main__":
    import argparse
    import uvicorn
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8001)
    args = parser.parse_args()
    uvicorn.run(app, host="0.0.0.0", port=args.port)