rm-background / app.py
magnusdtd's picture
Add initial FastAPI application with Docker support and file upload functionality
33f9dbc
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse, HTMLResponse
from rembg import remove
import uvicorn
from PIL import Image
from io import BytesIO
app = FastAPI()
@app.get("/", response_class=HTMLResponse)
async def read_index():
with open("templates/index.html", "r") as file:
return file.read()
@app.post("/")
async def upload_file(file: UploadFile = File(...)):
if not file:
raise HTTPException(status_code=400, detail="No file uploaded")
if file.filename == '':
raise HTTPException(status_code=400, detail="No file selected")
input_image = Image.open(file.file)
output_image = remove(input_image, post_process_mask=True)
img_io = BytesIO()
output_image.save(img_io, 'PNG')
img_io.seek(0)
return StreamingResponse(img_io, media_type="image/png", headers={"Content-Disposition": f"attachment; filename={file.filename}_rmng.png"})
if __name__ == '__main__':
uvicorn.run(app, host='0.0.0.0', port=5100, debug=True)