| from fastapi import FastAPI |
| from pydantic import BaseModel |
| from transparent_background import Remover |
| from PIL import Image |
| import io |
| import requests |
| import base64 |
| import warnings |
|
|
| warnings.filterwarnings("ignore") |
|
|
| app = FastAPI(title="Background Remover API") |
|
|
| |
| remover = Remover(mode="base-nightly") |
|
|
|
|
| class RemoveBackgroundRequest(BaseModel): |
| image_url: str |
| sod_type: str = "rgba" |
|
|
|
|
| class RemoveBackgroundResponse(BaseModel): |
| image_base64: str |
|
|
|
|
| @app.get("/") |
| def welcome(): |
| return {"message": "welcome"} |
|
|
|
|
| @app.post("/remove-background", response_model=RemoveBackgroundResponse) |
| def remove_background(payload: RemoveBackgroundRequest): |
| |
| response = requests.get(payload.image_url, timeout=10) |
| response.raise_for_status() |
| pil_image = Image.open(io.BytesIO(response.content)).convert("RGB") |
|
|
| |
| result = remover.process(pil_image, type=payload.sod_type) |
|
|
| |
| buf = io.BytesIO() |
| result.save(buf, format="PNG") |
| buf.seek(0) |
|
|
| |
| img_base64 = base64.b64encode(buf.getvalue()).decode("utf-8") |
|
|
| return {"image_base64": img_base64} |
|
|
|
|
| if __name__ == "__main__": |
| import uvicorn |
| uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True) |
|
|