from fastapi import FastAPI, UploadFile, File
from fastapi.responses import Response
import uvicorn
import shutil
import uuid
import os

from model_loader import model_manager
from onnx_processor import create_background_remover

app = FastAPI() 
# Chargement du modèle MODNet
modnet_path = model_manager.get_model_path("modnet")
processor = create_background_remover(model_type="modnet", model_path=modnet_path)

@app.post("/remove-bg/")
async def remove_bg(file: UploadFile = File(...)):
    # Sauvegarde temporaire
    input_path = f"/tmp/{uuid.uuid4()}.png"
    output_path = f"/tmp/{uuid.uuid4()}_output.png"

    with open(input_path, "wb") as f:
        shutil.copyfileobj(file.file, f)

    # Traitement de l'image
    try:
        processor.remove_background(image_path=input_path, output_path=output_path)
        with open(output_path, "rb") as out_file:
            return Response(content=out_file.read(), media_type="image/png")
    except Exception as e:
        return {"error": str(e)}
    finally:
        # Nettoyage
        if os.path.exists(input_path):
            os.remove(input_path)
        if os.path.exists(output_path):
            os.remove(output_path)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)