Update rembg/bg.py

rembg/bg.py

@@ -24,24 +24,19 @@ from .sessions.base import BaseSession
 
 kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
 
+
 class ReturnType(Enum):
     BYTES = 0
     PILLOW = 1
     NDARRAY = 2
 
-# Argumentos configurables desde la interfaz Gradio
-# Estos valores predeterminados pueden modificarse en la interfaz.
-foreground_threshold = 240
-background_threshold = 10
-erode_structure_size = 10
-alpha_matting = False
-only_mask = False
-post_process_mask = False
-bgcolor = None
 
 def alpha_matting_cutout(
     img: PILImage,
     mask: PILImage,
+    foreground_threshold: int,
+    background_threshold: int,
+    erode_structure_size: int,
 ) -> PILImage:
     if img.mode == "RGBA" or img.mode == "CMYK":
         img = img.convert("RGB")
@@ -77,15 +72,67 @@ def alpha_matting_cutout(
 
     return cutout
 
+
+def naive_cutout(img: PILImage, mask: PILImage) -> PILImage:
+    empty = Image.new("RGBA", (img.size), 0)
+    cutout = Image.composite(img, empty, mask)
+    return cutout
+
+
+def get_concat_v_multi(imgs: List[PILImage]) -> PILImage:
+    pivot = imgs.pop(0)
+    for im in imgs:
+        pivot = get_concat_v(pivot, im)
+    return pivot
+
+
+def get_concat_v(img1: PILImage, img2: PILImage) -> PILImage:
+    dst = Image.new("RGBA", (img1.width, img1.height + img2.height))
+    dst.paste(img1, (0, 0))
+    dst.paste(img2, (0, img1.height))
+    return dst
+
+
+def post_process(mask: np.ndarray) -> np.ndarray:
+    """
+    Post Process the mask for a smooth boundary by applying Morphological Operations
+    Research based on paper: https://www.sciencedirect.com/science/article/pii/S2352914821000757
+    args:
+        mask: Binary Numpy Mask
+    """
+    mask = morphologyEx(mask, MORPH_OPEN, kernel)
+    mask = GaussianBlur(mask, (5, 5), sigmaX=2, sigmaY=2, borderType=BORDER_DEFAULT)
+    mask = np.where(mask < 127, 0, 255).astype(np.uint8)  # convert again to binary
+    return mask
+
+
+def apply_background_color(img: PILImage, color: Tuple[int, int, int, int]) -> PILImage:
+    r, g, b, a = color
+    colored_image = Image.new("RGBA", img.size, (r, g, b, a))
+    colored_image.paste(img, mask=img)
+
+    return colored_image
+
+
+def fix_image_orientation(img: PILImage) -> PILImage:
+    return ImageOps.exif_transpose(img)
+
+
+def download_models() -> None:
+    for session in sessions_class:
+        session.download_models()
+
+
 def remove(
     data: Union[bytes, PILImage, np.ndarray],
-    alpha_matting: bool = alpha_matting,
-    alpha_matting_foreground_threshold: int = foreground_threshold,
-    alpha_matting_background_threshold: int = background_threshold,
-    alpha_matting_erode_size: int = erode_structure_size,
-    only_mask: bool = only_mask,
-    post_process_mask: bool = post_process_mask,
-    bgcolor: Optional[Tuple[int, int, int, int]] = bgcolor,
+    alpha_matting: bool = False,
+    alpha_matting_foreground_threshold: int = 240,
+    alpha_matting_background_threshold: int = 10,
+    alpha_matting_erode_size: int = 10,
+    session: Optional[BaseSession] = None,
+    only_mask: bool = False,
+    post_process_mask: bool = False,
+    bgcolor: Optional[Tuple[int, int, int, int]] = None,
     *args: Optional[Any],
     **kwargs: Optional[Any]
 ) -> Union[bytes, PILImage, np.ndarray]:
@@ -99,7 +146,7 @@ def remove(
         return_type = ReturnType.NDARRAY
         img = Image.fromarray(data)
     else:
-        raise ValueError("Input type {} is not supported.".format(type(data))
+        raise ValueError("Input type {} is not supported.".format(type(data)))
 
     # Fix image orientation
     img = fix_image_orientation(img)
@@ -122,6 +169,9 @@ def remove(
                 cutout = alpha_matting_cutout(
                     img,
                     mask,
+                    alpha_matting_foreground_threshold,
+                    alpha_matting_background_threshold,
+                    alpha_matting_erode_size,
                 )
             except ValueError:
                 cutout = naive_cutout(img, mask)