rembg/__init__.py

@@ -1,11 +1,6 @@
-import sys
-import warnings
-
-if not (sys.version_info.major == 3 and sys.version_info.minor == 9):
-    warnings.warn("This library is only for Python 3.9", RuntimeWarning)
-
 from . import _version
 
 __version__ = _version.get_versions()["version"]
 
 from .bg import remove
+from .session_factory import new_session

rembg/_version.py

@@ -24,8 +24,8 @@ def get_keywords():
     # each be defined on a line of their own. _version.py will just call
     # get_keywords().
     git_refnames = " (HEAD -> main)"
-    git_full = "3bc1c1af99ebd47dd08d02763fc754d70d42afea"
-    git_date = "2022-06-16 23:00:14 -0300"
+    git_full = "edc9fe27dff030cf6c2f29ef9a66c32d6e3f4658"
+    git_date = "2022-11-28 08:14:19 -0300"
     keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
     return keywords
 

rembg/bg.py

@@ -3,16 +3,26 @@ from enum import Enum
 from typing import List, Optional, Union
 
 import numpy as np
+from cv2 import (
+    BORDER_DEFAULT,
+    MORPH_ELLIPSE,
+    MORPH_OPEN,
+    GaussianBlur,
+    getStructuringElement,
+    morphologyEx,
+)
 from PIL import Image
 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.morphology import binary_erosion
+from scipy.ndimage import binary_erosion
 
 from .session_base import BaseSession
 from .session_factory import new_session
 
+kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
+
 
 class ReturnType(Enum):
     BYTES = 0
@@ -27,6 +37,10 @@ def alpha_matting_cutout(
     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)
 
@@ -79,6 +93,19 @@ def get_concat_v(img1: PILImage, img2: PILImage) -> PILImage:
     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 remove(
     data: Union[bytes, PILImage, np.ndarray],
     alpha_matting: bool = False,
@@ -87,6 +114,7 @@ def remove(
     alpha_matting_erode_size: int = 10,
     session: Optional[BaseSession] = None,
     only_mask: bool = False,
+    post_process_mask: bool = False,
 ) -> Union[bytes, PILImage, np.ndarray]:
 
     if isinstance(data, PILImage):
@@ -108,6 +136,9 @@ def remove(
     cutouts = []
 
     for mask in masks:
+        if post_process_mask:
+            mask = Image.fromarray(post_process(np.array(mask)))
+
         if only_mask:
             cutout = mask
 

rembg/cli.py

@@ -33,7 +33,9 @@ def main() -> None:
     "-m",
     "--model",
     default="u2net",
-    type=click.Choice(["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg"]),
+    type=click.Choice(
+        ["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta"]
+    ),
     show_default=True,
     show_choices=True,
     help="model name",
@@ -76,6 +78,13 @@ def main() -> None:
     show_default=True,
     help="output only the mask",
 )
+@click.option(
+    "-ppm",
+    "--post-process-mask",
+    is_flag=True,
+    show_default=True,
+    help="post process the mask",
+)
 @click.argument(
     "input", default=(None if sys.stdin.isatty() else "-"), type=click.File("rb")
 )
@@ -93,7 +102,9 @@ def i(model: str, input: IO, output: IO, **kwargs) -> None:
     "-m",
     "--model",
     default="u2net",
-    type=click.Choice(["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg"]),
+    type=click.Choice(
+        ["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta"]
+    ),
     show_default=True,
     show_choices=True,
     help="model name",
@@ -136,6 +147,13 @@ def i(model: str, input: IO, output: IO, **kwargs) -> None:
     show_default=True,
     help="output only the mask",
 )
+@click.option(
+    "-ppm",
+    "--post-process-mask",
+    is_flag=True,
+    show_default=True,
+    help="post process the mask",
+)
 @click.option(
     "-w",
     "--watch",
@@ -243,7 +261,15 @@ def p(
     show_default=True,
     help="log level",
 )
-def s(port: int, log_level: str) -> None:
+@click.option(
+    "-t",
+    "--threads",
+    default=None,
+    type=int,
+    show_default=True,
+    help="number of worker threads",
+)
+def s(port: int, log_level: str, threads: int) -> None:
     sessions: dict[str, BaseSession] = {}
     tags_metadata = [
         {
@@ -284,6 +310,7 @@ def s(port: int, log_level: str) -> None:
         u2netp = "u2netp"
         u2net_human_seg = "u2net_human_seg"
         u2net_cloth_seg = "u2net_cloth_seg"
+        silueta = "silueta"
 
     class CommonQueryParams:
         def __init__(
@@ -309,6 +336,7 @@ def s(port: int, log_level: str) -> None:
                 default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
             ),
             om: bool = Query(default=False, description="Only Mask"),
+            ppm: bool = Query(default=False, description="Post Process Mask"),
         ):
             self.model = model
             self.a = a
@@ -316,6 +344,7 @@ def s(port: int, log_level: str) -> None:
             self.ab = ab
             self.ae = ae
             self.om = om
+            self.ppm = ppm
 
     class CommonQueryPostParams:
         def __init__(
@@ -341,6 +370,7 @@ def s(port: int, log_level: str) -> None:
                 default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
             ),
             om: bool = Form(default=False, description="Only Mask"),
+            ppm: bool = Form(default=False, description="Post Process Mask"),
         ):
             self.model = model
             self.a = a
@@ -348,6 +378,7 @@ def s(port: int, log_level: str) -> None:
             self.ab = ab
             self.ae = ae
             self.om = om
+            self.ppm = ppm
 
     def im_without_bg(content: bytes, commons: CommonQueryParams) -> Response:
         return Response(
@@ -361,10 +392,19 @@ def s(port: int, log_level: str) -> None:
                 alpha_matting_background_threshold=commons.ab,
                 alpha_matting_erode_size=commons.ae,
                 only_mask=commons.om,
+                post_process_mask=commons.ppm,
             ),
             media_type="image/png",
         )
 
+    @app.on_event("startup")
+    def startup():
+        if threads is not None:
+            from anyio import CapacityLimiter
+            from anyio.lowlevel import RunVar
+
+            RunVar("_default_thread_limiter").set(CapacityLimiter(threads))
+
     @app.get(
         path="/",
         tags=["Background Removal"],

rembg/session_base.py

@@ -18,7 +18,7 @@ class BaseSession:
         std: Tuple[float, float, float],
         size: Tuple[int, int],
     ) -> Dict[str, np.ndarray]:
-        im = img.convert("RGB").resize(size, Image.LANCZOS)
+        im = img.convert("RGB").resize(size, Image.Resampling.LANCZOS)
 
         im_ary = np.array(im)
         im_ary = im_ary / np.max(im_ary)

rembg/session_factory.py

@@ -5,50 +5,56 @@ from contextlib import redirect_stdout
 from pathlib import Path
 from typing import Type
 
-import gdown
 import onnxruntime as ort
+import pooch
 
 from .session_base import BaseSession
 from .session_cloth import ClothSession
 from .session_simple import SimpleSession
 
 
-def new_session(model_name: str) -> BaseSession:
+def new_session(model_name: str = "u2net") -> BaseSession:
     session_class: Type[BaseSession]
+    md5 = "60024c5c889badc19c04ad937298a77b"
+    url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
+    session_class = SimpleSession
 
     if model_name == "u2netp":
         md5 = "8e83ca70e441ab06c318d82300c84806"
-        url = "https://drive.google.com/uc?id=1tNuFmLv0TSNDjYIkjEdeH1IWKQdUA4HR"
-        session_class = SimpleSession
-    elif model_name == "u2net":
-        md5 = "60024c5c889badc19c04ad937298a77b"
-        url = "https://drive.google.com/uc?id=1tCU5MM1LhRgGou5OpmpjBQbSrYIUoYab"
+        url = (
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
+        )
         session_class = SimpleSession
     elif model_name == "u2net_human_seg":
         md5 = "c09ddc2e0104f800e3e1bb4652583d1f"
-        url = "https://drive.google.com/uc?id=1ZfqwVxu-1XWC1xU1GHIP-FM_Knd_AX5j"
+        url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx"
         session_class = SimpleSession
     elif model_name == "u2net_cloth_seg":
         md5 = "2434d1f3cb744e0e49386c906e5a08bb"
-        url = "https://drive.google.com/uc?id=15rKbQSXQzrKCQurUjZFg8HqzZad8bcyz"
+        url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx"
         session_class = ClothSession
-    else:
-        assert AssertionError(
-            "Choose between u2net, u2netp, u2net_human_seg or u2net_cloth_seg"
+    elif model_name == "silueta":
+        md5 = "55e59e0d8062d2f5d013f4725ee84782"
+        url = (
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx"
         )
+        session_class = SimpleSession
 
-    home = os.getenv("U2NET_HOME", os.path.join("~", ".u2net"))
-    path = Path(home).expanduser() / f"{model_name}.onnx"
-    path.parents[0].mkdir(parents=True, exist_ok=True)
+    u2net_home = os.getenv(
+        "U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net")
+    )
+
+    fname = f"{model_name}.onnx"
+    path = Path(u2net_home).expanduser()
+    full_path = Path(u2net_home).expanduser() / fname
 
-    if not path.exists():
-        with redirect_stdout(sys.stderr):
-            gdown.download(url, str(path), use_cookies=False)
-    else:
-        hashing = hashlib.new("md5", path.read_bytes(), usedforsecurity=False)
-        if hashing.hexdigest() != md5:
-            with redirect_stdout(sys.stderr):
-                gdown.download(url, str(path), use_cookies=False)
+    pooch.retrieve(
+        url,
+        f"md5:{md5}",
+        fname=fname,
+        path=Path(u2net_home).expanduser(),
+        progressbar=True,
+    )
 
     sess_opts = ort.SessionOptions()
 
@@ -58,6 +64,8 @@ def new_session(model_name: str) -> BaseSession:
     return session_class(
         model_name,
         ort.InferenceSession(
-            str(path), providers=ort.get_available_providers(), sess_options=sess_opts
+            str(full_path),
+            providers=ort.get_available_providers(),
+            sess_options=sess_opts,
         ),
     )

rembg/session_simple.py

@@ -25,6 +25,6 @@ class SimpleSession(BaseSession):
         pred = np.squeeze(pred)
 
         mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
-        mask = mask.resize(img.size, Image.LANCZOS)
+        mask = mask.resize(img.size, Image.Resampling.LANCZOS)
 
         return [mask]