File size: 1,394 Bytes
6d97ee1
 
b6019a9
f50d763
 
6cba480
f50d763
 
6d97ee1
6d31469
f50d763
8369e7c
6d97ee1
b8a4f46
 
 
 
 
f0e753b
6d97ee1
 
 
 
 
 
 
 
 
 
 
 
 
8058e4c
f50d763
 
 
8058e4c
6d97ee1
 
 
 
f50d763
6d97ee1
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
from fastapi import FastAPI, Request, UploadFile, File
from fastapi.responses import StreamingResponse
from fastapi.responses import JSONResponse
from rembg import remove
from PIL import Image
import datetime
import io

app = FastAPI()
app.router.redirect_slashes = False

@app.api_route("/", methods=["GET", "HEAD"])
async def index(request: Request):
    return JSONResponse({
        "status": "ok",
        "message": "Server is alive",
        "timestamp": datetime.datetime.utcnow().isoformat() + "Z"
    })


@app.post('/remove-bg')
async def remove_bg(image: UploadFile = File(...)):
    """Accepts an uploaded image file, removes its background and returns a PNG.

    Input: multipart form file field named `image`.
    Output: PNG image stream with Content-Disposition attachment `no-bg.png`.
    """
    contents = await image.read()
    input_image = Image.open(io.BytesIO(contents))

    # perform background removal (rembg accepts a PIL Image)
    output = remove(input_image)

    buffer = io.BytesIO()
    output.save(buffer, format='PNG')
    buffer.seek(0)

    headers = {"Content-Disposition": "attachment; filename=\"no-bg.png\""}
    return StreamingResponse(buffer, media_type='image/png', headers=headers)

# If you want to run locally with `python api.py` for quick testing
if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=7860)