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

app = FastAPI()

# Limita o número de threads do onnxruntime
options = ort.SessionOptions()
options.intra_op_num_threads = 2

def resize_image(image, max_size=1024):
    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)  # Usando LANCZOS em vez de ANTIALIAS
    return image

def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
    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 e ajusta brilho/contraste
        image = resize_image(image, max_size=1024)
        image = adjust_brightness_contrast(image)
        
        # Remove o fundo da imagem usando rembg
        output = rembg.remove(image, model="u2netp", 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))