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, ImageFilter
import rembg
import onnxruntime as ort

app = FastAPI()

# Configurações do onnxruntime para CPU
options = ort.SessionOptions()
options.intra_op_num_threads = 2  # Limita o número de threads para evitar sobrecarga
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL  # Execução sequencial para melhor desempenho em CPU

def resize_image(image, max_size=512):
    """Redimensiona a imagem para uma largura máxima de 512px, mantendo a proporção."""
    width, height = image.size
    if width > max_size or height > max_size:
        ratio = min(max_size / width, max_size / height)
        new_size = (int(width * ratio), int(height * ratio))
        image = image.resize(new_size, Image.Resampling.LANCZOS)
    return image

@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))
        
        # Pré-processamento: apenas redimensiona para 512px
        image = resize_image(image, max_size=512)
        
        # Remove o fundo da imagem usando rembg
        output = rembg.remove(image, session_options=options)
        
        # Pós-processamento: suaviza as bordas
        output = output.filter(ImageFilter.SMOOTH_MORE)
        
        # 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))