app.py

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, ImageEnhance
import rembg
import onnxruntime as ort
from functools import lru_cache

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

# Cache do modelo u2net (carrega apenas uma vez)
@lru_cache(maxsize=1)
def load_u2net_model():
    return rembg.get_model("u2net")

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

def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
    """Ajusta o brilho e o contraste da imagem."""
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(contrast)
    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: redimensiona para 512px e ajusta brilho/contraste
        image = resize_image(image, max_size=512)
        image = adjust_brightness_contrast(image)
        
        # Remove o fundo da imagem usando rembg
        model = load_u2net_model()  # Carrega o modelo (usando cache)
        output = rembg.remove(image, model=model, 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))