Update app.py

app.py

@@ -8,6 +8,8 @@ import requests
 from io import BytesIO
 from PIL import Image, ImageFilter, ImageEnhance
 import rembg
+import cv2
+import numpy as np
 import onnxruntime as ort
 
 app = FastAPI()
@@ -16,14 +18,16 @@ app = FastAPI()
 options = ort.SessionOptions()
 options.intra_op_num_threads = 2
 
+# Função para redimensionar a imagem
 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
+        image = image.resize(new_size, Image.Resampling.LANCZOS)  # Usando LANCZOS para redimensionamento
     return image
 
+# Função para ajustar brilho e contraste
 def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
     enhancer = ImageEnhance.Brightness(image)
     image = enhancer.enhance(brightness)
@@ -31,6 +35,29 @@ def adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
     image = enhancer.enhance(contrast)
     return image
 
+# Função para remover sombras
+def remove_shadows(image):
+    # Converte a imagem para um array numpy (compatível com OpenCV)
+    image_np = np.array(image)
+    
+    # Converte a imagem para escala de cinza
+    gray = cv2.cvtColor(image_np, cv2.COLOR_RGBA2GRAY)
+    
+    # Aplica limiarização para remover sombras
+    _, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)
+    
+    # Aplica operações morfológicas para melhorar a máscara
+    kernel = np.ones((3, 3), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)  # Remove ruídos
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # Fecha buracos
+    
+    # Aplica a máscara na imagem original
+    result = cv2.bitwise_and(image_np, image_np, mask=mask)
+    
+    # Converte o resultado de volta para uma imagem PIL
+    return Image.fromarray(result)
+
+# Endpoint da API
 @app.get("/remove-background")
 async def remove_background(image_url: str):
     try:
@@ -46,7 +73,10 @@ async def remove_background(image_url: str):
         image = adjust_brightness_contrast(image)
         
         # Remove o fundo da imagem usando rembg
-        output = rembg.remove(image, model="u2netp", session_options=options)
+        output = rembg.remove(image, model="u2net", session_options=options)
+        
+        # Remove sombras
+        output = remove_shadows(output)
         
         # Pós-processamento: suaviza as bordas
         output = output.filter(ImageFilter.SMOOTH_MORE)