app.py

import os
import base64
import shutil
import tempfile
from gradio import Server
from gradio.data_classes import FileData
from gradio_client import Client, handle_file
from fastapi import UploadFile, File
from fastapi.responses import HTMLResponse, FileResponse

app = Server()

bg_client = Client("ysharma/background-removal-copy")


@app.api(name="remove_background")
def remove_background(image_path: FileData) -> FileData:
    """Remove background from an image. Returns path to transparent PNG."""
    result = bg_client.predict(f=handle_file(image_path["path"]), api_name="/png")
    return FileData(path=result)


@app.post("/api/remove-background")
async def remove_bg_endpoint(file: UploadFile = File(...)):
    """FastAPI endpoint for the custom frontend to call."""
    suffix = os.path.splitext(file.filename or ".png")[1]
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        shutil.copyfileobj(file.file, tmp)
        tmp_path = tmp.name
    try:
        result_path = bg_client.predict(
            f=handle_file(tmp_path), api_name="/png"
        )
        with open(result_path, "rb") as f:
            fg_b64 = base64.b64encode(f.read()).decode()
        return {"foreground": f"data:image/png;base64,{fg_b64}"}
    finally:
        os.unlink(tmp_path)


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


# Server-side pixel-perfect export via Playwright
from export_server import register_export_endpoint
register_export_endpoint(app)


@app.get("/export_client.js")
async def serve_export_js():
    return FileResponse(
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "export_client.js"),
        media_type="application/javascript",
    )


app.launch(show_error=True)