Spaces:
Sleeping
Sleeping
File size: 3,105 Bytes
add1294 afac93b add1294 4b2ecc5 d77e4f1 e8b21d9 d77e4f1 96eebff 0e69a3b 72a2ce8 4b2ecc5 39f2554 1aab113 39f2554 1aab113 72a2ce8 0e69a3b 72a2ce8 39f2554 1aab113 96eebff 1aab113 0e69a3b 72a2ce8 0e69a3b d77e4f1 4b2ecc5 72a2ce8 d77e4f1 72a2ce8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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
from concurrent.futures import ThreadPoolExecutor
import asyncio
import gc
from asyncio import Semaphore
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
# Cria uma sessão global do rembg para reutilização
session = rembg.new_session(model_name="u2net", session_options=options)
# Pool de threads para executar tarefas bloqueantes
executor = ThreadPoolExecutor(max_workers=4)
# Semáforo para limitar requisições simultâneas
semaphore = Semaphore(4) # Ajuste conforme a capacidade da sua CPU/RAM
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 process_image(image_url):
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 (reutiliza a sessão global)
output = rembg.remove(image, session=session)
# 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)
# Libera memória explicitamente
del image, output
gc.collect() # Força a liberação de memória
return img_byte_arr
except Exception as e:
raise e
@app.get("/remove-background")
async def remove_background(image_url: str):
async with semaphore: # Limita o número de requisições simultâneas
try:
# Executa o processamento da imagem em um thread separado
loop = asyncio.get_event_loop()
img_byte_arr = await loop.run_in_executor(executor, process_image, image_url)
# 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)) |