Spaces:
Sleeping
Sleeping
import os | |
os.environ["NUMBA_CACHE_DIR"] = "/tmp/numba_cache" | |
os.environ["U2NET_HOME"] = "/tmp/u2net" | |
from fastapi import FastAPI, HTTPException | |
from pydantic import BaseModel | |
from io import BytesIO | |
from PIL import Image, ImageFilter | |
import rembg | |
import onnxruntime as ort | |
from concurrent.futures import ThreadPoolExecutor | |
import asyncio | |
import gc | |
import base64 | |
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 | |
# Pool de threads para executar tarefas bloqueantes | |
executor = ThreadPoolExecutor(max_workers=4) | |
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_data): | |
# Decodifica a imagem base64 | |
image = Image.open(BytesIO(base64.b64decode(image_data))) | |
# 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) | |
# Codifica a imagem processada em base64 | |
return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8') | |
class ImageBase64(BaseModel): | |
image: str # A imagem em base64 | |
async def remove_background(image_data: ImageBase64): | |
try: | |
# Executa o processamento da imagem em um thread separado | |
loop = asyncio.get_event_loop() | |
processed_image_base64 = await loop.run_in_executor(executor, process_image, image_data.image) | |
# Retorna a imagem processada em base64 | |
return {"processed_image": processed_image_base64} | |
except Exception as e: | |
raise HTTPException(status_code=400, detail=str(e)) | |
finally: | |
# Força a coleta de lixo após cada requisição | |
gc.collect() |