Update app.py

app.py

@@ -9,16 +9,22 @@ from io import BytesIO
 from PIL import Image, ImageFilter, ImageEnhance
 import rembg
 import onnxruntime as ort
-import cv2
-import numpy as np
+from functools import lru_cache
 
 app = FastAPI()
 
-# Limita o número de threads do onnxruntime
+# Configurações do onnxruntime para CPU
 options = ort.SessionOptions()
-options.intra_op_num_threads = 2
+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=1024):
+# 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)
@@ -27,23 +33,13 @@ def resize_image(image, max_size=1024):
     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
 
-def remove_shadows(image):
-    # Converte a imagem para escala de cinza
-    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2GRAY)
-    
-    # Aplica limiarização para remover áreas escuras (sombras)
-    _, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)
-    
-    # Aplica a máscara na imagem original
-    result = cv2.bitwise_and(np.array(image), np.array(image), mask=mask)
-    return Image.fromarray(result)
-
 @app.get("/remove-background")
 async def remove_background(image_url: str):
     try:
@@ -54,17 +50,15 @@ async def remove_background(image_url: str):
         # 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)
+        # 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
-        output = rembg.remove(image, model="u2net_human_seg", session_options=options)
-        
-        # Remove sombras
-        output = remove_shadows(output)
+        model = load_u2net_model()  # Carrega o modelo (usando cache)
+        output = rembg.remove(image, model=model, session_options=options)
         
-        # Suaviza a imagem
+        # Pós-processamento: suaviza as bordas
         output = output.filter(ImageFilter.SMOOTH_MORE)
         
         # Converte a imagem de volta para bytes