"""Nukki Studio — FastAPI 서버.""" import base64 import io from pathlib import Path from fastapi import FastAPI, File, Form, UploadFile from fastapi.responses import HTMLResponse, JSONResponse, Response from fastapi.staticfiles import StaticFiles from PIL import Image from . import config, pipeline from .device import device_name from .registry import REGISTRY WEB = Path(__file__).parent / "web" app = FastAPI(title="Nukki Studio") @app.get("/", response_class=HTMLResponse) def index(): return (WEB / "index.html").read_text(encoding="utf-8") @app.get("/api/health") def health(): return { "device": device_name(), "status": REGISTRY.status, "seg_models": list(config.SEG_MODELS.keys()), "modes": list(config.MODES.keys()), "default_mode": config.DEFAULT_MODE, } @app.post("/api/warmup") def warmup(seg_model: str = Form("general"), matte: bool = Form(True)): REGISTRY.segmenter(seg_model) if matte: REGISTRY.matte() return {"ok": True, "loaded": REGISTRY.status["loaded"]} def _b64(image: Image.Image) -> str: return "data:image/png;base64," + base64.b64encode(pipeline.encode_png(image)).decode() @app.post("/api/remove") async def remove( file: UploadFile = File(...), mode: str = Form(config.DEFAULT_MODE), subject: str = Form("auto"), seg_model: str = Form(""), matte_size: str = Form(config.DEFAULT_VITMATTE), feather: float = Form(0.0), edge_shift: float = Form(0.0), band_scale: float = Form(1.0), max_side: int = Form(0), matte_max_side: int = Form(0), want_mask: bool = Form(True), ): raw = await file.read() try: image = Image.open(io.BytesIO(raw)) image.load() except Exception as e: return JSONResponse({"error": f"이미지를 열 수 없습니다: {e}"}, status_code=400) rgba, info = pipeline.process( image, mode=mode, subject=subject, seg_model=seg_model or None, matte_size=matte_size, feather=feather, edge_shift=edge_shift, band_scale=band_scale, max_side=max_side, matte_max_side=matte_max_side, ) payload = {"result": _b64(rgba), "info": info} if want_mask: mask = rgba.getchannel("A").convert("L") payload["mask"] = _b64(mask.convert("RGB")) return JSONResponse(payload) @app.post("/api/remove.png") async def remove_png(file: UploadFile = File(...), mode: str = Form(config.DEFAULT_MODE), subject: str = Form("auto")): """순수 PNG 바이너리 반환 (CLI/자동화용).""" raw = await file.read() image = Image.open(io.BytesIO(raw)) rgba, _ = pipeline.process(image, mode=mode, subject=subject) return Response(content=pipeline.encode_png(rgba), media_type="image/png") # 정적 자산 app.mount("/static", StaticFiles(directory=str(WEB)), name="static")