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)