Update app.py

app.py

@@ -8,9 +8,9 @@ import requests
 from io import BytesIO
 from PIL import Image, ImageFilter, ImageEnhance
 import rembg
+import onnxruntime as ort
 import cv2
 import numpy as np
-import onnxruntime as ort
 
 app = FastAPI()
 
@@ -18,16 +18,14 @@ 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 para redimensionamento
+        image = image.resize(new_size, Image.Resampling.LANCZOS)
     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)
@@ -35,29 +33,17 @@ 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)
+    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2GRAY)
     
-    # Aplica limiarização para remover sombras
+    # Aplica limiarização para remover áreas escuras (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
+    result = cv2.bitwise_and(np.array(image), np.array(image), mask=mask)
     return Image.fromarray(result)
 
-# Endpoint da API
 @app.get("/remove-background")
 async def remove_background(image_url: str):
     try:
@@ -73,12 +59,12 @@ async def remove_background(image_url: str):
         image = adjust_brightness_contrast(image)
         
         # Remove o fundo da imagem usando rembg
-        output = rembg.remove(image, model="u2net", session_options=options)
+        output = rembg.remove(image, model="u2net_human_seg", session_options=options)
         
         # Remove sombras
         output = remove_shadows(output)
         
-        # Pós-processamento: suaviza as bordas
+        # Suaviza a imagem
         output = output.filter(ImageFilter.SMOOTH_MORE)
         
         # Converte a imagem de volta para bytes