import os
os.environ["NUMBA_CACHE_DIR"] = "/tmp/numba_cache"
os.environ["U2NET_HOME"] = "/tmp/u2net"

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import requests
from io import BytesIO
from PIL import Image
import rembg

app = FastAPI()

@app.get("/remove-background")
async def remove_background(image_url: str):
    try:
        # Baixa a imagem da URL fornecida
        response = requests.get(image_url)
        response.raise_for_status()
        
        # Abre a imagem usando Pillow
        image = Image.open(BytesIO(response.content))
        
        # Remove o fundo da imagem usando rembg
        output = rembg.remove(image)
        
        # Converte a imagem de volta para bytes
        img_byte_arr = BytesIO()
        output.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        # Retorna a imagem processada diretamente no navegador
        return StreamingResponse(img_byte_arr, media_type="image/png")
    
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))