115 lines
3.2 KiB
Python
115 lines
3.2 KiB
Python
|
"""FastAPI backend for garment photo cleaning."""
|
||
|
|
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
from fastapi import FastAPI, File, Form, HTTPException, UploadFile, status
|
||
|
|
from fastapi.middleware.cors import CORSMiddleware
|
||
|
|
from fastapi.responses import FileResponse, JSONResponse
|
||
|
|
from fastapi.staticfiles import StaticFiles
|
||
|
|
|
||
|
|
from cleaner import (
|
||
|
|
DEFAULT_MODEL,
|
||
|
|
CleanerError,
|
||
|
|
_get_image_info,
|
||
|
|
build_prompt,
|
||
|
|
clean_garment,
|
||
|
|
encode_image_bytes,
|
||
|
|
)
|
||
|
|
|
||
|
|
DEFAULT_PROMPT = build_prompt("beige")
|
||
|
|
STATIC_DIR = Path(__file__).resolve().parent / "static"
|
||
|
|
STATIC_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
app = FastAPI(title="unwrap-clothes")
|
||
|
|
|
||
|
|
app.add_middleware(
|
||
|
|
CORSMiddleware,
|
||
|
|
allow_origins=["*"],
|
||
|
|
allow_credentials=True,
|
||
|
|
allow_methods=["*"],
|
||
|
|
allow_headers=["*"],
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@app.get("/api/health")
|
||
|
|
def health_check():
|
||
|
|
return {"status": "ok"}
|
||
|
|
|
||
|
|
|
||
|
|
@app.get("/")
|
||
|
|
def read_root():
|
||
|
|
index_file = STATIC_DIR / "index.html"
|
||
|
|
if index_file.is_file():
|
||
|
|
return FileResponse(index_file)
|
||
|
|
return JSONResponse({"status": "ok", "message": "unwrap-clothes backend is running"})
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
@app.post("/api/clean")
|
||
|
|
async def clean_endpoint(
|
||
|
|
file: UploadFile = File(...),
|
||
|
|
prompt: str = Form(default=DEFAULT_PROMPT),
|
||
|
|
model: str = Form(default=DEFAULT_MODEL),
|
||
|
|
api_key: str | None = Form(default=None),
|
||
|
|
restore_res: bool = Form(default=True),
|
||
|
|
):
|
||
|
|
effective_api_key = (api_key.strip() if api_key and api_key.strip() else None) or os.environ.get(
|
||
|
|
"OPENROUTER_API_KEY"
|
||
|
|
)
|
||
|
|
if not effective_api_key:
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
||
|
|
detail="API key is required. Provide it in the form or set OPENROUTER_API_KEY.",
|
||
|
|
)
|
||
|
|
|
||
|
|
file_bytes = await file.read()
|
||
|
|
if not file_bytes:
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
||
|
|
detail="Uploaded file is empty.",
|
||
|
|
)
|
||
|
|
|
||
|
|
orig_dims, _ = _get_image_info(file_bytes)
|
||
|
|
if orig_dims is None:
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
||
|
|
detail="Invalid image file. Could not decode image.",
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
clean_result = clean_garment(
|
||
|
|
image_data=file_bytes,
|
||
|
|
prompt=prompt,
|
||
|
|
api_key=effective_api_key,
|
||
|
|
model=model,
|
||
|
|
restore_res=restore_res,
|
||
|
|
)
|
||
|
|
except CleanerError as exc:
|
||
|
|
raise HTTPException(
|
||
|
|
status_code=status.HTTP_502_BAD_GATEWAY,
|
||
|
|
detail=str(exc),
|
||
|
|
)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"image": encode_image_bytes(clean_result.image_bytes, mime=clean_result.media_type),
|
||
|
|
"cost": clean_result.cost,
|
||
|
|
"width": clean_result.width,
|
||
|
|
"height": clean_result.height,
|
||
|
|
"original_dimensions": list(clean_result.original_dimensions)
|
||
|
|
if clean_result.original_dimensions
|
||
|
|
else None,
|
||
|
|
"model_dimensions": list(clean_result.model_dimensions)
|
||
|
|
if clean_result.model_dimensions
|
||
|
|
else None,
|
||
|
|
"was_rescaled": clean_result.was_rescaled,
|
||
|
|
"media_type": clean_result.media_type,
|
||
|
|
}
|
||
|
|
|
||
|
|
if STATIC_DIR.is_dir():
|
||
|
|
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
import uvicorn
|
||
|
|
|
||
|
|
uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True)
|