Update handler.py

handler.py

@@ -1,48 +1,45 @@
-# These HF deployment codes refer to https://huggingface.co/not-lain/BiRefNet/raw/main/handler.py.
-from typing import Dict, List, Any, Tuple
+from typing import Dict, Any, Tuple
 import os
 import requests
 from io import BytesIO
-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 ---
+# --- 1. Configuración ---
 torch.set_float32_matmul_precision(["high", "highest"][0])
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-class ImagePreprocessor():
-    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
-        self.transform_image = transforms.Compose([
-            transforms.Resize(resolution),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-        ])
-    def proc(self, image: Image.Image) -> torch.Tensor:
-        image = self.transform_image(image)
-        return image
-
 usage_to_weights_file = {
     'General': 'BiRefNet',
-    'General-HR': 'BiRefNet_HR',
     'General-Lite': 'BiRefNet_lite',
     'General-Lite-2K': 'BiRefNet_lite-2K',
     'General-reso_512': 'BiRefNet-reso_512',
-    'General-legacy': 'BiRefNet-legacy'
+    'General-HR': 'BiRefNet_HR'
 }
 
 usage = 'General'
 resolution = (1024, 1024)
 half_precision = True
 
+class ImagePreprocessor():
+    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
+        self.transform_image = transforms.Compose([
+            transforms.Resize(resolution),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+    def proc(self, image: Image.Image) -> torch.Tensor:
+        image = self.transform_image(image)
+        return image
+
 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
+            '/'.join(('zhengpeng7', usage_to_weights_file[usage])), 
+            trust_remote_code=True
         )
         self.birefnet.to(device)
         self.birefnet.eval()
@@ -50,12 +47,11 @@ class EndpointHandler():
             self.birefnet.half()
 
     def __call__(self, data: Dict[str, Any]):
-        # ---------------------------------------------------------
-        # 1. LÓGICA BLINDADA DE ENTRADA (Mantenemos lo que ya funcionaba)
-        # ---------------------------------------------------------
+        # --- PASO 1: Carga Segura de la Imagen ---
         image_src = data["inputs"]
         image_ori = None
-
+        
+        # Detectamos qué nos enviaron (Objeto, URL o Bytes)
         if hasattr(image_src, 'convert') or isinstance(image_src, Image.Image):
             image_ori = image_src
         elif isinstance(image_src, str):
@@ -72,13 +68,11 @@ class EndpointHandler():
                     image_ori = Image.fromarray(image_src)
                 except Exception:
                     image_ori = image_src
-        
-        # Convertimos a RGB para asegurar consistencia
+
+        # Convertimos a RGB (Esto limpia cualquier rareza del archivo original y asegura color)
         image = image_ori.convert('RGB')
         
-        # ---------------------------------------------------------
-        # 2. INFERENCIA (Detectar qué es fondo y qué es producto)
-        # ---------------------------------------------------------
+        # --- PASO 2: La IA detecta la silueta ---
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image)
         image_proc = image_proc.unsqueeze(0)
@@ -88,31 +82,15 @@ class EndpointHandler():
         
         pred = preds[0].squeeze()
         
-        # ---------------------------------------------------------
-        # 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)
+        # --- PASO 3: Recorte Limpio (Sin matemáticas raras) ---
+        # Convertimos la predicción en una máscara (imagen en blanco y negro)
+        mask_pil = transforms.ToPILImage()(pred)
         
-        # 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]))
+        # Redimensionamos la máscara al tamaño EXACTO de la foto original
+        mask_pil = mask_pil.resize(image.size, resample=Image.Resampling.LANCZOS)
         
-        # 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)
+        # ✨ MAGIA: Simplemente le decimos a la foto original "Usa esta transparencia"
+        # No tocamos los canales de color (RGB), solo añadimos el canal Alpha.
+        image.putalpha(mask_pil)
         
-        return final_image
+        return image