Update rembg/bg.py

rembg/bg.py

@@ -1,26 +1,142 @@
 import io
-from typing import Any, Union
-import numpy as np
-from PIL import Image
 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 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))
 
-from rembg import new_session, remove
-from rembg.sessions.base import BaseSession
-from rembg.util.util import fix_image_orientation
 
 class ReturnType(Enum):
     BYTES = 0
     PILLOW = 1
     NDARRAY = 2
 
-def remove_background(
-    data: Union[bytes, Image.Image, np.ndarray],
-    alpha_influence: float = 0.5,
-    segmentation_strength: float = 0.5,
-    smoothing: float = 0.5,
-    model: str = "u2net",
-) -> Union[bytes, Image.Image, np.ndarray]:
-    if isinstance(data, Image.Image):
+
+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")
+
+    img = np.asarray(img)
+    mask = np.asarray(mask)
+
+    is_foreground = mask > foreground_threshold
+    is_background = mask < background_threshold
+
+    structure = None
+    if erode_structure_size > 0:
+        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)
+
+    trimap = np.full(mask.shape, dtype=np.uint8, fill_value=128)
+    trimap[is_foreground] = 255
+    trimap[is_background] = 0
+
+    img_normalized = img / 255.0
+    trimap_normalized = trimap / 255.0
+
+    alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
+    foreground = estimate_foreground_ml(img_normalized, alpha)
+    cutout = stack_images(foreground, alpha)
+
+    cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)
+    cutout = Image.fromarray(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 = 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]:
+    if isinstance(data, PILImage):
         return_type = ReturnType.PILLOW
         img = data
     elif isinstance(data, bytes):
@@ -30,27 +146,56 @@ def remove_background(
         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)
-    session = new_session(model)
-    output = remove(
-        img,
-        alpha_matting=True,
-        alpha_matting_foreground_threshold=alpha_influence * 255,
-        alpha_matting_background_threshold=(1 - alpha_influence) * 255,
-        alpha_matting_erode_size=int(segmentation_strength * 20),
-        alpha_matting_smoothing=smoothing,
-        session=session
-    )
-
-    if return_type == ReturnType.PILLOW:
-        return output
-    elif return_type == ReturnType.NDARRAY:
-        return np.array(output)
-    else:
-        bio = io.BytesIO()
-        output.save(bio, "PNG")
-        bio.seek(0)
-        return bio.read()
 
+    if session is None:
+        session = new_session("u2net", *args, **kwargs)
+
+    masks = session.predict(img, *args, **kwargs)
+    cutouts = []
+
+    for mask in masks:
+        if post_process_mask:
+            mask = Image.fromarray(post_process(np.array(mask)))
+
+        if only_mask:
+            cutout = mask
+
+        elif alpha_matting:
+            try:
+                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)
+
+        else:
+            cutout = naive_cutout(img, mask)
+
+        cutouts.append(cutout)
+
+    cutout = img
+    if len(cutouts) > 0:
+        cutout = get_concat_v_multi(cutouts)
+
+    if bgcolor is not None and not only_mask:
+        cutout = apply_background_color(cutout, bgcolor)
+
+    if ReturnType.PILLOW == return_type:
+        return cutout
+
+    if ReturnType.NDARRAY == return_type:
+        return np.asarray(cutout)
+
+    bio = io.BytesIO()
+    cutout.save(bio, "PNG")
+    bio.seek(0)
+
+    return bio.read()