Update handler.py

handler.py

@@ -3,46 +3,17 @@ from typing import Dict, List, Any, Tuple
 import os
 import requests
 from io import BytesIO
-import cv2
-import numpy as np
+import cv2  # Importante para el redimensionado preciso
+import numpy as np # Importante para la manipulación de píxeles
 from PIL import Image
 import torch
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
 
+# --- Configuración Básica ---
 torch.set_float32_matmul_precision(["high", "highest"][0])
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-### image_proc.py
-def refine_foreground(image, mask, r=90):
-    if mask.size != image.size:
-        mask = mask.resize(image.size)
-    image = np.array(image) / 255.0
-    mask = np.array(mask) / 255.0
-    estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r)
-    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
-    return image_masked
-
-def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
-    alpha = alpha[:, :, None]
-    F, blur_B = FB_blur_fusion_foreground_estimator(image, image, image, alpha, r)
-    return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
-
-def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
-    # Detección segura para helpers internos
-    if hasattr(image, 'size') or isinstance(image, Image.Image):
-        image = np.array(image) / 255.0
-    
-    blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None]
-    blurred_FA = cv2.blur(F * alpha, (r, r))
-    blurred_F = blurred_FA / (blurred_alpha + 1e-5)
-    blurred_B1A = cv2.blur(B * (1 - alpha), (r, r))
-    blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
-    F = blurred_F + alpha * \
-        (image - alpha * blurred_F - (1 - alpha) * blurred_B)
-    F = np.clip(F, 0, 1)
-    return F, blurred_B
-
 class ImagePreprocessor():
     def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
         self.transform_image = transforms.Compose([
@@ -60,13 +31,6 @@ usage_to_weights_file = {
     'General-Lite': 'BiRefNet_lite',
     'General-Lite-2K': 'BiRefNet_lite-2K',
     'General-reso_512': 'BiRefNet-reso_512',
-    'Matting': 'BiRefNet-matting',
-    'Matting-HR': 'BiRefNet_HR-Matting',
-    'Portrait': 'BiRefNet-portrait',
-    'DIS': 'BiRefNet-DIS5K',
-    'HRSOD': 'BiRefNet-HRSOD',
-    'COD': 'BiRefNet-COD',
-    'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
     'General-legacy': 'BiRefNet-legacy'
 }
 
@@ -76,6 +40,7 @@ half_precision = True
 
 class EndpointHandler():
     def __init__(self, path=''):
+        # Carga del modelo
         self.birefnet = AutoModelForImageSegmentation.from_pretrained(
             '/'.join(('zhengpeng7', usage_to_weights_file[usage])), trust_remote_code=True
         )
@@ -85,48 +50,69 @@ class EndpointHandler():
             self.birefnet.half()
 
     def __call__(self, data: Dict[str, Any]):
-        print('data["inputs"] type:', type(data["inputs"])) # Log para debug
+        # ---------------------------------------------------------
+        # 1. LÓGICA BLINDADA DE ENTRADA (Mantenemos lo que ya funcionaba)
+        # ---------------------------------------------------------
         image_src = data["inputs"]
+        image_ori = None
 
-        # --- LOGICA BLINDADA ---
-        # 1. Si ya es una imagen (tiene atributo 'size' o 'convert'), úsala directo.
         if hasattr(image_src, 'convert') or isinstance(image_src, Image.Image):
             image_ori = image_src
-        
-        # 2. Si es una ruta de archivo o URL (String)
         elif isinstance(image_src, str):
             if os.path.isfile(image_src):
                 image_ori = Image.open(image_src)
             else:
                 response = requests.get(image_src)
-                image_data = BytesIO(response.content)
-                image_ori = Image.open(image_data)
-        
-        # 3. Último recurso: Bytes crudos o Arrays
+                image_ori = Image.open(BytesIO(response.content))
         else:
             try:
-                # Intenta abrirlo como bytes (lo más común si falla el paso 1)
                 image_ori = Image.open(BytesIO(image_src))
             except Exception:
                 try:
-                    # Intenta como array de numpy
                     image_ori = Image.fromarray(image_src)
                 except Exception:
-                    # Si falla todo, asume que YA es una imagen que falló la detección
                     image_ori = image_src
-        # -----------------------
-
+        
+        # Convertimos a RGB para asegurar consistencia
         image = image_ori.convert('RGB')
         
+        # ---------------------------------------------------------
+        # 2. INFERENCIA (Detectar qué es fondo y qué es producto)
+        # ---------------------------------------------------------
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image)
         image_proc = image_proc.unsqueeze(0)
         
         with torch.no_grad():
             preds = self.birefnet(image_proc.to(device).half() if half_precision else image_proc.to(device))[-1].sigmoid().cpu()
+        
         pred = preds[0].squeeze()
         
-        pred_pil = transforms.ToPILImage()(pred)
-        image_masked = refine_foreground(image, pred_pil)
-        image_masked.putalpha(pred_pil.resize(image.size))
-        return image_masked
+        # ---------------------------------------------------------
+        # 3. RECONSTRUCCIÓN MATEMÁTICA (Solución al problema del negro)
+        # ---------------------------------------------------------
+        
+        # A. Convertimos la predicción a array numpy y normalizamos
+        mask_np = pred.numpy()
+        mask_np = (mask_np - mask_np.min()) / (mask_np.max() - mask_np.min() + 1e-8)
+        
+        # B. Convertimos la imagen original a matriz de números [Alto, Ancho, 3]
+        image_np = np.array(image)
+        
+        # C. Redimensionamos la máscara al tamaño EXACTO de la imagen original
+        # (Esto evita deformaciones o bordes extraños)
+        mask_resized = cv2.resize(mask_np, (image_np.shape[1], image_np.shape[0]))
+        
+        # D. Creamos una imagen vacía de 4 canales (RGBA - Rojo, Verde, Azul, Alfa)
+        rgba_image = np.zeros((image_np.shape[0], image_np.shape[1], 4), dtype=np.uint8)
+        
+        # E. Copiamos los colores ORIGINALES (Sin modificarlos ni mezclarlos)
+        rgba_image[:, :, :3] = image_np
+        
+        # F. Aplicamos la máscara al canal Alfa (Transparencia)
+        rgba_image[:, :, 3] = (mask_resized * 255).astype(np.uint8)
+
+        # G. Convertimos de vuelta a imagen PIL para devolverla
+        final_image = Image.fromarray(rgba_image)
+        
+        return final_image