Update rembg/bg.py

rembg/bg.py

@@ -1,41 +1,27 @@
 import io
-from enum import Enum
-from typing import Any, List, Optional, Tuple, Union
-
 import numpy as np
-from cv2 import (
-    BORDER_DEFAULT,
-    MORPH_ELLIPSE,
-    MORPH_OPEN,
-    GaussianBlur,
-    getStructuringElement,
-    morphologyEx,
-)
 from PIL import Image, ImageOps
-from PIL.Image import Image as PILImage
+from typing import Any, List, Optional, Tuple, Union
+from enum import Enum
 from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
 from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
-from pymatting.util.util import stack_images
 from scipy.ndimage import binary_erosion
 
-from .session_factory import new_session
-from .sessions import sessions_class
-from .sessions.base import BaseSession
-
-kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
-
+# Definir los valores por defecto para los parámetros
+DEFAULT_ALPHA_INFLUENCE = 0.5
+DEFAULT_SEGMENTATION_STRENGTH = 0.5
+DEFAULT_SMOOTHING = 0.5
 
 class ReturnType(Enum):
     BYTES = 0
     PILLOW = 1
     NDARRAY = 2
 
-
 def alpha_matting_cutout(
     img: PILImage,
     mask: PILImage,
-    foreground_threshold: int,
-    background_threshold: int,
+    alpha_influence: float,
+    segmentation_strength: float,
     erode_structure_size: int,
 ) -> PILImage:
     if img.mode == "RGBA" or img.mode == "CMYK":
@@ -44,14 +30,12 @@ def alpha_matting_cutout(
     img = np.asarray(img)
     mask = np.asarray(mask)
 
-    is_foreground = mask > foreground_threshold
-    is_background = mask < background_threshold
+    is_foreground = mask > alpha_influence
+    is_background = mask < (1.0 - alpha_influence)
 
     structure = None
     if erode_structure_size > 0:
-        structure = np.ones(
-            (erode_structure_size, erode_structure_size), dtype=np.uint8
-        )
+        structure = np.ones((erode_structure_size, erode_structure_size), dtype=np.uint8)
 
     is_foreground = binary_erosion(is_foreground, structure=structure)
     is_background = binary_erosion(is_background, structure=structure, border_value=1)
@@ -71,8 +55,7 @@ def alpha_matting_cutout(
     cutout = Image.fromarray(cutout)
 
     return cutout
-
-
+#aca termina la modificacion
 def naive_cutout(img: PILImage, mask: PILImage) -> PILImage:
     empty = Image.new("RGBA", (img.size), 0)
     cutout = Image.composite(img, empty, mask)