Commit changes to main

handler.py

@@ -11,7 +11,6 @@ from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 ### image_proc.py
@@ -24,22 +23,18 @@ def refine_foreground(image, mask, r=90):
     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):
     # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation
     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):
     if 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 * \
@@ -47,7 +42,6 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
     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([
@@ -55,7 +49,6 @@ class ImagePreprocessor():
             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
@@ -111,7 +104,13 @@ class EndpointHandler():
         """
         print('data["inputs"] = ', data["inputs"])
         image_src = data["inputs"]
-        if isinstance(image_src, str):
+
+        # ------------------------------------------------------------------
+        # MODIFICACION REPUESTOS MOM: Soporte para imágenes directas (Bytes/PIL)
+        # ------------------------------------------------------------------
+        if isinstance(image_src, Image.Image):
+            image_ori = image_src
+        elif isinstance(image_src, str):
             if os.path.isfile(image_src):
                 image_ori = Image.open(image_src)
             else:
@@ -119,21 +118,32 @@ class EndpointHandler():
                 image_data = BytesIO(response.content)
                 image_ori = Image.open(image_data)
         else:
-            image_ori = Image.fromarray(image_src)
-
+            try:
+                # Intento leer como array (comportamiento original)
+                image_ori = Image.fromarray(image_src)
+            except Exception:
+                # Fallback: Intento leer como bytes crudos (para Odoo)
+                try:
+                    image_ori = Image.open(BytesIO(image_src))
+                except Exception:
+                    # Si falla, intentamos array de nuevo como último recurso
+                    image_ori = Image.fromarray(image_src)
+        # ------------------------------------------------------------------
+        
         image = image_ori.convert('RGB')
+        
         # Preprocess the image
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image)
         image_proc = image_proc.unsqueeze(0)
-
+        
         # Prediction
         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()
-
+        
         # Show Results
         pred_pil = transforms.ToPILImage()(pred)
         image_masked = refine_foreground(image, pred_pil)
         image_masked.putalpha(pred_pil.resize(image.size))
-        return image_masked
+        return image_masked