app.py

@@ -9,9 +9,7 @@ def inference(file, af, mask, model):
     im = cv2.imread(file, cv2.IMREAD_COLOR)
     cv2.imwrite(os.path.join("input.png"), im)
 
-    from rembg import remove
-    from rembg.session_base import BaseSession
-    from rembg.session_factory import new_session
+    from rembg import new_session, remove
 
     input_path = 'input.png'
     output_path = 'output.png'
@@ -19,15 +17,15 @@ def inference(file, af, mask, model):
     with open(input_path, 'rb') as i:
         with open(output_path, 'wb') as o:
             input = i.read()
-            sessions: dict[str, BaseSession] = {}
             output = remove(
                 input, 
-                session=sessions.setdefault(
-                    model, new_session(model)
-                ),
+                session = new_session(model), 
                 alpha_matting_erode_size = af, 
                 only_mask = (True if mask == "Mask only" else False)
-            )
+            ) 
+            
+            
+            
             o.write(output)
     return os.path.join("output.png")
         
@@ -40,7 +38,7 @@ gr.Interface(
     inference, 
     [
         gr.inputs.Image(type="filepath", label="Input"),
-        gr.inputs.Slider(10, 25, default=10, label="Alpha matting"), 
+        gr.inputs.Slider(10, 25, default=10, label="Alpha matting erode size"), 
         gr.inputs.Radio(
             [
                 "Default", 
@@ -55,14 +53,16 @@ gr.Interface(
             "u2netp", 
             "u2net_human_seg", 
             "u2net_cloth_seg", 
-            "silueta"
+            "silueta",
+            "isnet-general-use",
+            "sam",
             ], 
             type="value",
             default="u2net",
             label="Models"
         ),
     ], 
-    gr.outputs.Image(type="file", label="Output"),
+    gr.outputs.Image(type="filepath", label="Output"),
     title=title,
     description=description,
     article=article,

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 = "d62227d5866e2178e88f06074917484a4424082e"
-    git_date = "2022-12-10 11:51:49 -0300"
+    git_full = "e47b2a0ed405a5a30f42bacb142b107f7a4b6536"
+    git_date = "2023-04-26 20:40:21 -0300"
     keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
     return keywords
 

rembg/bg.py

@@ -1,6 +1,6 @@
 import io
 from enum import Enum
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
 from cv2 import (
@@ -18,8 +18,8 @@ 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_base import BaseSession
 from .session_factory import new_session
+from .sessions.base import BaseSession
 
 kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
 
@@ -37,7 +37,6 @@ def alpha_matting_cutout(
     background_threshold: int,
     erode_structure_size: int,
 ) -> PILImage:
-
     if img.mode == "RGBA" or img.mode == "CMYK":
         img = img.convert("RGB")
 
@@ -106,6 +105,14 @@ def post_process(mask: np.ndarray) -> np.ndarray:
     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 remove(
     data: Union[bytes, PILImage, np.ndarray],
     alpha_matting: bool = False,
@@ -115,8 +122,10 @@ def remove(
     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
@@ -130,9 +139,9 @@ def remove(
         raise ValueError("Input type {} is not supported.".format(type(data)))
 
     if session is None:
-        session = new_session("u2net")
+        session = new_session("u2net", *args, **kwargs)
 
-    masks = session.predict(img)
+    masks = session.predict(img, *args, **kwargs)
     cutouts = []
 
     for mask in masks:
@@ -163,6 +172,9 @@ def remove(
     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
 

rembg/cli.py

@@ -1,25 +1,7 @@
-import pathlib
-import sys
-import time
-from enum import Enum
-from typing import IO, cast
-
-import aiohttp
 import click
-import filetype
-import uvicorn
-from asyncer import asyncify
-from fastapi import Depends, FastAPI, File, Form, Query
-from fastapi.middleware.cors import CORSMiddleware
-from starlette.responses import Response
-from tqdm import tqdm
-from watchdog.events import FileSystemEvent, FileSystemEventHandler
-from watchdog.observers import Observer
 
 from . import _version
-from .bg import remove
-from .session_base import BaseSession
-from .session_factory import new_session
+from .commands import command_functions
 
 
 @click.group()
@@ -28,413 +10,5 @@ def main() -> None:
     pass
 
 
-@main.command(help="for a file as input")
-@click.option(
-    "-m",
-    "--model",
-    default="u2net",
-    type=click.Choice(
-        ["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta"]
-    ),
-    show_default=True,
-    show_choices=True,
-    help="model name",
-)
-@click.option(
-    "-a",
-    "--alpha-matting",
-    is_flag=True,
-    show_default=True,
-    help="use alpha matting",
-)
-@click.option(
-    "-af",
-    "--alpha-matting-foreground-threshold",
-    default=240,
-    type=int,
-    show_default=True,
-    help="trimap fg threshold",
-)
-@click.option(
-    "-ab",
-    "--alpha-matting-background-threshold",
-    default=10,
-    type=int,
-    show_default=True,
-    help="trimap bg threshold",
-)
-@click.option(
-    "-ae",
-    "--alpha-matting-erode-size",
-    default=10,
-    type=int,
-    show_default=True,
-    help="erode size",
-)
-@click.option(
-    "-om",
-    "--only-mask",
-    is_flag=True,
-    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")
-)
-@click.argument(
-    "output",
-    default=(None if sys.stdin.isatty() else "-"),
-    type=click.File("wb", lazy=True),
-)
-def i(model: str, input: IO, output: IO, **kwargs) -> None:
-    output.write(remove(input.read(), session=new_session(model), **kwargs))
-
-
-@main.command(help="for a folder as input")
-@click.option(
-    "-m",
-    "--model",
-    default="u2net",
-    type=click.Choice(
-        ["u2net", "u2netp", "u2net_human_seg", "u2net_cloth_seg", "silueta"]
-    ),
-    show_default=True,
-    show_choices=True,
-    help="model name",
-)
-@click.option(
-    "-a",
-    "--alpha-matting",
-    is_flag=True,
-    show_default=True,
-    help="use alpha matting",
-)
-@click.option(
-    "-af",
-    "--alpha-matting-foreground-threshold",
-    default=240,
-    type=int,
-    show_default=True,
-    help="trimap fg threshold",
-)
-@click.option(
-    "-ab",
-    "--alpha-matting-background-threshold",
-    default=10,
-    type=int,
-    show_default=True,
-    help="trimap bg threshold",
-)
-@click.option(
-    "-ae",
-    "--alpha-matting-erode-size",
-    default=10,
-    type=int,
-    show_default=True,
-    help="erode size",
-)
-@click.option(
-    "-om",
-    "--only-mask",
-    is_flag=True,
-    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",
-    default=False,
-    is_flag=True,
-    show_default=True,
-    help="watches a folder for changes",
-)
-@click.argument(
-    "input",
-    type=click.Path(
-        exists=True,
-        path_type=pathlib.Path,
-        file_okay=False,
-        dir_okay=True,
-        readable=True,
-    ),
-)
-@click.argument(
-    "output",
-    type=click.Path(
-        exists=False,
-        path_type=pathlib.Path,
-        file_okay=False,
-        dir_okay=True,
-        writable=True,
-    ),
-)
-def p(
-    model: str, input: pathlib.Path, output: pathlib.Path, watch: bool, **kwargs
-) -> None:
-    session = new_session(model)
-
-    def process(each_input: pathlib.Path) -> None:
-        try:
-            mimetype = filetype.guess(each_input)
-            if mimetype is None:
-                return
-            if mimetype.mime.find("image") < 0:
-                return
-
-            each_output = (output / each_input.name).with_suffix(".png")
-            each_output.parents[0].mkdir(parents=True, exist_ok=True)
-
-            if not each_output.exists():
-                each_output.write_bytes(
-                    cast(
-                        bytes,
-                        remove(each_input.read_bytes(), session=session, **kwargs),
-                    )
-                )
-
-                if watch:
-                    print(
-                        f"processed: {each_input.absolute()} -> {each_output.absolute()}"
-                    )
-        except Exception as e:
-            print(e)
-
-    inputs = list(input.glob("**/*"))
-    if not watch:
-        inputs = tqdm(inputs)
-
-    for each_input in inputs:
-        if not each_input.is_dir():
-            process(each_input)
-
-    if watch:
-        observer = Observer()
-
-        class EventHandler(FileSystemEventHandler):
-            def on_any_event(self, event: FileSystemEvent) -> None:
-                if not (
-                    event.is_directory or event.event_type in ["deleted", "closed"]
-                ):
-                    process(pathlib.Path(event.src_path))
-
-        event_handler = EventHandler()
-        observer.schedule(event_handler, input, recursive=False)
-        observer.start()
-
-        try:
-            while True:
-                time.sleep(1)
-
-        finally:
-            observer.stop()
-            observer.join()
-
-
-@main.command(help="for a http server")
-@click.option(
-    "-p",
-    "--port",
-    default=5000,
-    type=int,
-    show_default=True,
-    help="port",
-)
-@click.option(
-    "-l",
-    "--log_level",
-    default="info",
-    type=str,
-    show_default=True,
-    help="log level",
-)
-@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 = [
-        {
-            "name": "Background Removal",
-            "description": "Endpoints that perform background removal with different image sources.",
-            "externalDocs": {
-                "description": "GitHub Source",
-                "url": "https://github.com/danielgatis/rembg",
-            },
-        },
-    ]
-    app = FastAPI(
-        title="Rembg",
-        description="Rembg is a tool to remove images background. That is it.",
-        version=_version.get_versions()["version"],
-        contact={
-            "name": "Daniel Gatis",
-            "url": "https://github.com/danielgatis",
-            "email": "[email protected]",
-        },
-        license_info={
-            "name": "MIT License",
-            "url": "https://github.com/danielgatis/rembg/blob/main/LICENSE.txt",
-        },
-        openapi_tags=tags_metadata,
-    )
-
-    app.add_middleware(
-        CORSMiddleware,
-        allow_credentials=True,
-        allow_origins=["*"],
-        allow_methods=["*"],
-        allow_headers=["*"],
-    )
-
-    class ModelType(str, Enum):
-        u2net = "u2net"
-        u2netp = "u2netp"
-        u2net_human_seg = "u2net_human_seg"
-        u2net_cloth_seg = "u2net_cloth_seg"
-        silueta = "silueta"
-
-    class CommonQueryParams:
-        def __init__(
-            self,
-            model: ModelType = Query(
-                default=ModelType.u2net,
-                description="Model to use when processing image",
-            ),
-            a: bool = Query(default=False, description="Enable Alpha Matting"),
-            af: int = Query(
-                default=240,
-                ge=0,
-                le=255,
-                description="Alpha Matting (Foreground Threshold)",
-            ),
-            ab: int = Query(
-                default=10,
-                ge=0,
-                le=255,
-                description="Alpha Matting (Background Threshold)",
-            ),
-            ae: int = Query(
-                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
-            self.af = af
-            self.ab = ab
-            self.ae = ae
-            self.om = om
-            self.ppm = ppm
-
-    class CommonQueryPostParams:
-        def __init__(
-            self,
-            model: ModelType = Form(
-                default=ModelType.u2net,
-                description="Model to use when processing image",
-            ),
-            a: bool = Form(default=False, description="Enable Alpha Matting"),
-            af: int = Form(
-                default=240,
-                ge=0,
-                le=255,
-                description="Alpha Matting (Foreground Threshold)",
-            ),
-            ab: int = Form(
-                default=10,
-                ge=0,
-                le=255,
-                description="Alpha Matting (Background Threshold)",
-            ),
-            ae: int = Form(
-                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
-            self.af = af
-            self.ab = ab
-            self.ae = ae
-            self.om = om
-            self.ppm = ppm
-
-    def im_without_bg(content: bytes, commons: CommonQueryParams) -> Response:
-        return Response(
-            remove(
-                content,
-                session=sessions.setdefault(
-                    commons.model.value, new_session(commons.model.value)
-                ),
-                alpha_matting=commons.a,
-                alpha_matting_foreground_threshold=commons.af,
-                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"],
-        summary="Remove from URL",
-        description="Removes the background from an image obtained by retrieving an URL.",
-    )
-    async def get_index(
-        url: str = Query(
-            default=..., description="URL of the image that has to be processed."
-        ),
-        commons: CommonQueryParams = Depends(),
-    ):
-        async with aiohttp.ClientSession() as session:
-            async with session.get(url) as response:
-                file = await response.read()
-                return await asyncify(im_without_bg)(file, commons)
-
-    @app.post(
-        path="/",
-        tags=["Background Removal"],
-        summary="Remove from Stream",
-        description="Removes the background from an image sent within the request itself.",
-    )
-    async def post_index(
-        file: bytes = File(
-            default=...,
-            description="Image file (byte stream) that has to be processed.",
-        ),
-        commons: CommonQueryPostParams = Depends(),
-    ):
-        return await asyncify(im_without_bg)(file, commons)
-
-    uvicorn.run(app, host="0.0.0.0", port=port, log_level=log_level)
+for command in command_functions:
+    main.add_command(command)

rembg/commands/__init__.py

@@ -0,0 +1,13 @@
+from importlib import import_module
+from pathlib import Path
+from pkgutil import iter_modules
+
+command_functions = []
+
+package_dir = Path(__file__).resolve().parent
+for _b, module_name, _p in iter_modules([str(package_dir)]):
+    module = import_module(f"{__name__}.{module_name}")
+    for attribute_name in dir(module):
+        attribute = getattr(module, attribute_name)
+        if attribute_name.endswith("_command"):
+            command_functions.append(attribute)

rembg/commands/i_command.py

@@ -0,0 +1,93 @@
+import json
+import sys
+from typing import IO
+
+import click
+
+from ..bg import remove
+from ..session_factory import new_session
+from ..sessions import sessions_names
+
+
+@click.command(
+    name="i",
+    help="for a file as input",
+)
+@click.option(
+    "-m",
+    "--model",
+    default="u2net",
+    type=click.Choice(sessions_names),
+    show_default=True,
+    show_choices=True,
+    help="model name",
+)
+@click.option(
+    "-a",
+    "--alpha-matting",
+    is_flag=True,
+    show_default=True,
+    help="use alpha matting",
+)
+@click.option(
+    "-af",
+    "--alpha-matting-foreground-threshold",
+    default=240,
+    type=int,
+    show_default=True,
+    help="trimap fg threshold",
+)
+@click.option(
+    "-ab",
+    "--alpha-matting-background-threshold",
+    default=10,
+    type=int,
+    show_default=True,
+    help="trimap bg threshold",
+)
+@click.option(
+    "-ae",
+    "--alpha-matting-erode-size",
+    default=10,
+    type=int,
+    show_default=True,
+    help="erode size",
+)
+@click.option(
+    "-om",
+    "--only-mask",
+    is_flag=True,
+    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(
+    "-bgc",
+    "--bgcolor",
+    default=None,
+    type=(int, int, int, int),
+    nargs=4,
+    help="Background color (R G B A) to replace the removed background with",
+)
+@click.option("-x", "--extras", type=str)
+@click.argument(
+    "input", default=(None if sys.stdin.isatty() else "-"), type=click.File("rb")
+)
+@click.argument(
+    "output",
+    default=(None if sys.stdin.isatty() else "-"),
+    type=click.File("wb", lazy=True),
+)
+def i_command(model: str, extras: str, input: IO, output: IO, **kwargs) -> None:
+    try:
+        kwargs.update(json.loads(extras))
+    except Exception:
+        pass
+
+    output.write(remove(input.read(), session=new_session(model), **kwargs))

rembg/commands/p_command.py

@@ -0,0 +1,181 @@
+import json
+import pathlib
+import time
+from typing import cast
+
+import click
+import filetype
+from tqdm import tqdm
+from watchdog.events import FileSystemEvent, FileSystemEventHandler
+from watchdog.observers import Observer
+
+from ..bg import remove
+from ..session_factory import new_session
+from ..sessions import sessions_names
+
+
+@click.command(
+    name="p",
+    help="for a folder as input",
+)
+@click.option(
+    "-m",
+    "--model",
+    default="u2net",
+    type=click.Choice(sessions_names),
+    show_default=True,
+    show_choices=True,
+    help="model name",
+)
+@click.option(
+    "-a",
+    "--alpha-matting",
+    is_flag=True,
+    show_default=True,
+    help="use alpha matting",
+)
+@click.option(
+    "-af",
+    "--alpha-matting-foreground-threshold",
+    default=240,
+    type=int,
+    show_default=True,
+    help="trimap fg threshold",
+)
+@click.option(
+    "-ab",
+    "--alpha-matting-background-threshold",
+    default=10,
+    type=int,
+    show_default=True,
+    help="trimap bg threshold",
+)
+@click.option(
+    "-ae",
+    "--alpha-matting-erode-size",
+    default=10,
+    type=int,
+    show_default=True,
+    help="erode size",
+)
+@click.option(
+    "-om",
+    "--only-mask",
+    is_flag=True,
+    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",
+    default=False,
+    is_flag=True,
+    show_default=True,
+    help="watches a folder for changes",
+)
+@click.option(
+    "-bgc",
+    "--bgcolor",
+    default=None,
+    type=(int, int, int, int),
+    nargs=4,
+    help="Background color (R G B A) to replace the removed background with",
+)
+@click.option("-x", "--extras", type=str)
+@click.argument(
+    "input",
+    type=click.Path(
+        exists=True,
+        path_type=pathlib.Path,
+        file_okay=False,
+        dir_okay=True,
+        readable=True,
+    ),
+)
+@click.argument(
+    "output",
+    type=click.Path(
+        exists=False,
+        path_type=pathlib.Path,
+        file_okay=False,
+        dir_okay=True,
+        writable=True,
+    ),
+)
+def p_command(
+    model: str,
+    extras: str,
+    input: pathlib.Path,
+    output: pathlib.Path,
+    watch: bool,
+    **kwargs,
+) -> None:
+    try:
+        kwargs.update(json.loads(extras))
+    except Exception:
+        pass
+
+    session = new_session(model)
+
+    def process(each_input: pathlib.Path) -> None:
+        try:
+            mimetype = filetype.guess(each_input)
+            if mimetype is None:
+                return
+            if mimetype.mime.find("image") < 0:
+                return
+
+            each_output = (output / each_input.name).with_suffix(".png")
+            each_output.parents[0].mkdir(parents=True, exist_ok=True)
+
+            if not each_output.exists():
+                each_output.write_bytes(
+                    cast(
+                        bytes,
+                        remove(each_input.read_bytes(), session=session, **kwargs),
+                    )
+                )
+
+                if watch:
+                    print(
+                        f"processed: {each_input.absolute()} -> {each_output.absolute()}"
+                    )
+        except Exception as e:
+            print(e)
+
+    inputs = list(input.glob("**/*"))
+    if not watch:
+        inputs = tqdm(inputs)
+
+    for each_input in inputs:
+        if not each_input.is_dir():
+            process(each_input)
+
+    if watch:
+        observer = Observer()
+
+        class EventHandler(FileSystemEventHandler):
+            def on_any_event(self, event: FileSystemEvent) -> None:
+                if not (
+                    event.is_directory or event.event_type in ["deleted", "closed"]
+                ):
+                    process(pathlib.Path(event.src_path))
+
+        event_handler = EventHandler()
+        observer.schedule(event_handler, input, recursive=False)
+        observer.start()
+
+        try:
+            while True:
+                time.sleep(1)
+
+        finally:
+            observer.stop()
+            observer.join()

rembg/commands/s_command.py

@@ -0,0 +1,238 @@
+import json
+from typing import Annotated, Optional, Tuple, cast
+
+import aiohttp
+import click
+import uvicorn
+from asyncer import asyncify
+from fastapi import Depends, FastAPI, File, Form, Query
+from fastapi.middleware.cors import CORSMiddleware
+from starlette.responses import Response
+
+from .._version import get_versions
+from ..bg import remove
+from ..session_factory import new_session
+from ..sessions import sessions_names
+from ..sessions.base import BaseSession
+
+
+@click.command(
+    name="s",
+    help="for a http server",
+)
+@click.option(
+    "-p",
+    "--port",
+    default=5000,
+    type=int,
+    show_default=True,
+    help="port",
+)
+@click.option(
+    "-l",
+    "--log_level",
+    default="info",
+    type=str,
+    show_default=True,
+    help="log level",
+)
+@click.option(
+    "-t",
+    "--threads",
+    default=None,
+    type=int,
+    show_default=True,
+    help="number of worker threads",
+)
+def s_command(port: int, log_level: str, threads: int) -> None:
+    sessions: dict[str, BaseSession] = {}
+    tags_metadata = [
+        {
+            "name": "Background Removal",
+            "description": "Endpoints that perform background removal with different image sources.",
+            "externalDocs": {
+                "description": "GitHub Source",
+                "url": "https://github.com/danielgatis/rembg",
+            },
+        },
+    ]
+    app = FastAPI(
+        title="Rembg",
+        description="Rembg is a tool to remove images background. That is it.",
+        version=get_versions()["version"],
+        contact={
+            "name": "Daniel Gatis",
+            "url": "https://github.com/danielgatis",
+            "email": "[email protected]",
+        },
+        license_info={
+            "name": "MIT License",
+            "url": "https://github.com/danielgatis/rembg/blob/main/LICENSE.txt",
+        },
+        openapi_tags=tags_metadata,
+    )
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_credentials=True,
+        allow_origins=["*"],
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    class CommonQueryParams:
+        def __init__(
+            self,
+            model: Annotated[
+                str, Query(regex=r"(" + "|".join(sessions_names) + ")")
+            ] = Query(
+                description="Model to use when processing image",
+            ),
+            a: bool = Query(default=False, description="Enable Alpha Matting"),
+            af: int = Query(
+                default=240,
+                ge=0,
+                le=255,
+                description="Alpha Matting (Foreground Threshold)",
+            ),
+            ab: int = Query(
+                default=10,
+                ge=0,
+                le=255,
+                description="Alpha Matting (Background Threshold)",
+            ),
+            ae: int = Query(
+                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"),
+            bgc: Optional[str] = Query(default=None, description="Background Color"),
+            extras: Optional[str] = Query(
+                default=None, description="Extra parameters as JSON"
+            ),
+        ):
+            self.model = model
+            self.a = a
+            self.af = af
+            self.ab = ab
+            self.ae = ae
+            self.om = om
+            self.ppm = ppm
+            self.extras = extras
+            self.bgc = (
+                cast(Tuple[int, int, int, int], tuple(map(int, bgc.split(","))))
+                if bgc
+                else None
+            )
+
+    class CommonQueryPostParams:
+        def __init__(
+            self,
+            model: Annotated[
+                str, Form(regex=r"(" + "|".join(sessions_names) + ")")
+            ] = Form(
+                description="Model to use when processing image",
+            ),
+            a: bool = Form(default=False, description="Enable Alpha Matting"),
+            af: int = Form(
+                default=240,
+                ge=0,
+                le=255,
+                description="Alpha Matting (Foreground Threshold)",
+            ),
+            ab: int = Form(
+                default=10,
+                ge=0,
+                le=255,
+                description="Alpha Matting (Background Threshold)",
+            ),
+            ae: int = Form(
+                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"),
+            bgc: Optional[str] = Query(default=None, description="Background Color"),
+            extras: Optional[str] = Query(
+                default=None, description="Extra parameters as JSON"
+            ),
+        ):
+            self.model = model
+            self.a = a
+            self.af = af
+            self.ab = ab
+            self.ae = ae
+            self.om = om
+            self.ppm = ppm
+            self.extras = extras
+            self.bgc = (
+                cast(Tuple[int, int, int, int], tuple(map(int, bgc.split(","))))
+                if bgc
+                else None
+            )
+
+    def im_without_bg(content: bytes, commons: CommonQueryParams) -> Response:
+        kwargs = {}
+
+        if commons.extras:
+            try:
+                kwargs.update(json.loads(commons.extras))
+            except Exception:
+                pass
+
+        return Response(
+            remove(
+                content,
+                session=sessions.setdefault(commons.model, new_session(commons.model)),
+                alpha_matting=commons.a,
+                alpha_matting_foreground_threshold=commons.af,
+                alpha_matting_background_threshold=commons.ab,
+                alpha_matting_erode_size=commons.ae,
+                only_mask=commons.om,
+                post_process_mask=commons.ppm,
+                bgcolor=commons.bgc,
+                **kwargs
+            ),
+            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"],
+        summary="Remove from URL",
+        description="Removes the background from an image obtained by retrieving an URL.",
+    )
+    async def get_index(
+        url: str = Query(
+            default=..., description="URL of the image that has to be processed."
+        ),
+        commons: CommonQueryParams = Depends(),
+    ):
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url) as response:
+                file = await response.read()
+                return await asyncify(im_without_bg)(file, commons)
+
+    @app.post(
+        path="/",
+        tags=["Background Removal"],
+        summary="Remove from Stream",
+        description="Removes the background from an image sent within the request itself.",
+    )
+    async def post_index(
+        file: bytes = File(
+            default=...,
+            description="Image file (byte stream) that has to be processed.",
+        ),
+        commons: CommonQueryPostParams = Depends(),
+    ):
+        return await asyncify(im_without_bg)(file, commons)  # type: ignore
+
+    uvicorn.run(app, host="0.0.0.0", port=port, log_level=log_level)

rembg/session_factory.py

@@ -1,71 +1,24 @@
-import hashlib
 import os
-import sys
-from contextlib import redirect_stdout
-from pathlib import Path
 from typing import Type
 
 import onnxruntime as ort
-import pooch
 
-from .session_base import BaseSession
-from .session_cloth import ClothSession
-from .session_simple import SimpleSession
+from .sessions import sessions_class
+from .sessions.base import BaseSession
+from .sessions.u2net import U2netSession
 
 
-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
+def new_session(model_name: str = "u2net", *args, **kwargs) -> BaseSession:
+    session_class: Type[BaseSession] = U2netSession
 
-    if model_name == "u2netp":
-        md5 = "8e83ca70e441ab06c318d82300c84806"
-        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://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://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx"
-        session_class = ClothSession
-    elif model_name == "silueta":
-        md5 = "55e59e0d8062d2f5d013f4725ee84782"
-        url = (
-            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx"
-        )
-        session_class = SimpleSession
-
-    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
-
-    pooch.retrieve(
-        url,
-        f"md5:{md5}",
-        fname=fname,
-        path=Path(u2net_home).expanduser(),
-        progressbar=True,
-    )
+    for sc in sessions_class:
+        if sc.name() == model_name:
+            session_class = sc
+            break
 
     sess_opts = ort.SessionOptions()
 
     if "OMP_NUM_THREADS" in os.environ:
         sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
 
-    return session_class(
-        model_name,
-        ort.InferenceSession(
-            str(full_path),
-            providers=ort.get_available_providers(),
-            sess_options=sess_opts,
-        ),
-    )
+    return session_class(model_name, sess_opts, *args, **kwargs)

rembg/sessions/__init__.py

@@ -0,0 +1,22 @@
+from importlib import import_module
+from inspect import isclass
+from pathlib import Path
+from pkgutil import iter_modules
+
+from .base import BaseSession
+
+sessions_class = []
+sessions_names = []
+
+package_dir = Path(__file__).resolve().parent
+for _b, module_name, _p in iter_modules([str(package_dir)]):
+    module = import_module(f"{__name__}.{module_name}")
+    for attribute_name in dir(module):
+        attribute = getattr(module, attribute_name)
+        if (
+            isclass(attribute)
+            and issubclass(attribute, BaseSession)
+            and attribute != BaseSession
+        ):
+            sessions_class.append(attribute)
+            sessions_names.append(attribute.name())

rembg/sessions/base.py

@@ -0,0 +1,63 @@
+import os
+from typing import Dict, List, Tuple
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+
+class BaseSession:
+    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
+        self.model_name = model_name
+        self.inner_session = ort.InferenceSession(
+            str(self.__class__.download_models()),
+            providers=ort.get_available_providers(),
+            sess_options=sess_opts,
+        )
+
+    def normalize(
+        self,
+        img: PILImage,
+        mean: Tuple[float, float, float],
+        std: Tuple[float, float, float],
+        size: Tuple[int, int],
+        *args,
+        **kwargs
+    ) -> Dict[str, np.ndarray]:
+        im = img.convert("RGB").resize(size, Image.LANCZOS)
+
+        im_ary = np.array(im)
+        im_ary = im_ary / np.max(im_ary)
+
+        tmpImg = np.zeros((im_ary.shape[0], im_ary.shape[1], 3))
+        tmpImg[:, :, 0] = (im_ary[:, :, 0] - mean[0]) / std[0]
+        tmpImg[:, :, 1] = (im_ary[:, :, 1] - mean[1]) / std[1]
+        tmpImg[:, :, 2] = (im_ary[:, :, 2] - mean[2]) / std[2]
+
+        tmpImg = tmpImg.transpose((2, 0, 1))
+
+        return {
+            self.inner_session.get_inputs()[0]
+            .name: np.expand_dims(tmpImg, 0)
+            .astype(np.float32)
+        }
+
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        raise NotImplementedError
+
+    @classmethod
+    def u2net_home(cls, *args, **kwargs):
+        return os.path.expanduser(
+            os.getenv(
+                "U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net")
+            )
+        )
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        raise NotImplementedError
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        raise NotImplementedError

rembg/sessions/dis.py

@@ -0,0 +1,47 @@
+import os
+from typing import List
+
+import numpy as np
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class DisSession(BaseSession):
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(img, (0.485, 0.456, 0.406), (1.0, 1.0, 1.0), (1024, 1024)),
+        )
+
+        pred = ort_outs[0][:, 0, :, :]
+
+        ma = np.max(pred)
+        mi = np.min(pred)
+
+        pred = (pred - mi) / (ma - mi)
+        pred = np.squeeze(pred)
+
+        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+        mask = mask.resize(img.size, Image.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname = f"{cls.name()}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx",
+            "md5:fc16ebd8b0c10d971d3513d564d01e29",
+            fname=fname,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "isnet-general-use"

rembg/sessions/sam.py

@@ -0,0 +1,165 @@
+import os
+from typing import List
+
+import numpy as np
+import onnxruntime as ort
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int):
+    scale = long_side_length * 1.0 / max(oldh, oldw)
+    newh, neww = oldh * scale, oldw * scale
+    neww = int(neww + 0.5)
+    newh = int(newh + 0.5)
+    return (newh, neww)
+
+
+def apply_coords(coords: np.ndarray, original_size, target_length) -> np.ndarray:
+    old_h, old_w = original_size
+    new_h, new_w = get_preprocess_shape(
+        original_size[0], original_size[1], target_length
+    )
+    coords = coords.copy().astype(float)
+    coords[..., 0] = coords[..., 0] * (new_w / old_w)
+    coords[..., 1] = coords[..., 1] * (new_h / old_h)
+    return coords
+
+
+def resize_longes_side(img: PILImage, size=1024):
+    w, h = img.size
+    if h > w:
+        new_h, new_w = size, int(w * size / h)
+    else:
+        new_h, new_w = int(h * size / w), size
+
+    return img.resize((new_w, new_h))
+
+
+def pad_to_square(img: np.ndarray, size=1024):
+    h, w = img.shape[:2]
+    padh = size - h
+    padw = size - w
+    img = np.pad(img, ((0, padh), (0, padw), (0, 0)), mode="constant")
+    img = img.astype(np.float32)
+    return img
+
+
+class SamSession(BaseSession):
+    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
+        self.model_name = model_name
+        paths = self.__class__.download_models()
+        self.encoder = ort.InferenceSession(
+            str(paths[0]),
+            providers=ort.get_available_providers(),
+            sess_options=sess_opts,
+        )
+        self.decoder = ort.InferenceSession(
+            str(paths[1]),
+            providers=ort.get_available_providers(),
+            sess_options=sess_opts,
+        )
+
+    def normalize(
+        self,
+        img: np.ndarray,
+        mean=(123.675, 116.28, 103.53),
+        std=(58.395, 57.12, 57.375),
+        size=(1024, 1024),
+        *args,
+        **kwargs,
+    ):
+        pixel_mean = np.array([*mean]).reshape(1, 1, -1)
+        pixel_std = np.array([*std]).reshape(1, 1, -1)
+        x = (img - pixel_mean) / pixel_std
+        return x
+
+    def predict(
+        self,
+        img: PILImage,
+        *args,
+        **kwargs,
+    ) -> List[PILImage]:
+        # Preprocess image
+        image = resize_longes_side(img)
+        image = np.array(image)
+        image = self.normalize(image)
+        image = pad_to_square(image)
+
+        input_labels = kwargs.get("input_labels")
+        input_points = kwargs.get("input_points")
+
+        if input_labels is None:
+            raise ValueError("input_labels is required")
+        if input_points is None:
+            raise ValueError("input_points is required")
+
+        # Transpose
+        image = image.transpose(2, 0, 1)[None, :, :, :]
+        # Run encoder (Image embedding)
+        encoded = self.encoder.run(None, {"x": image})
+        image_embedding = encoded[0]
+
+        # Add a batch index, concatenate a padding point, and transform.
+        onnx_coord = np.concatenate([input_points, np.array([[0.0, 0.0]])], axis=0)[
+            None, :, :
+        ]
+        onnx_label = np.concatenate([input_labels, np.array([-1])], axis=0)[
+            None, :
+        ].astype(np.float32)
+        onnx_coord = apply_coords(onnx_coord, img.size[::1], 1024).astype(np.float32)
+
+        # Create an empty mask input and an indicator for no mask.
+        onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
+        onnx_has_mask_input = np.zeros(1, dtype=np.float32)
+
+        decoder_inputs = {
+            "image_embeddings": image_embedding,
+            "point_coords": onnx_coord,
+            "point_labels": onnx_label,
+            "mask_input": onnx_mask_input,
+            "has_mask_input": onnx_has_mask_input,
+            "orig_im_size": np.array(img.size[::-1], dtype=np.float32),
+        }
+
+        masks, _, low_res_logits = self.decoder.run(None, decoder_inputs)
+        masks = masks > 0.0
+        masks = [
+            Image.fromarray((masks[i, 0] * 255).astype(np.uint8))
+            for i in range(masks.shape[0])
+        ]
+
+        return masks
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname_encoder = f"{cls.name()}_encoder.onnx"
+        fname_decoder = f"{cls.name()}_decoder.onnx"
+
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-encoder-quant.onnx",
+            "md5:13d97c5c79ab13ef86d67cbde5f1b250",
+            fname=fname_encoder,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-decoder-quant.onnx",
+            "md5:fa3d1c36a3187d3de1c8deebf33dd127",
+            fname=fname_decoder,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        return (
+            os.path.join(cls.u2net_home(), fname_encoder),
+            os.path.join(cls.u2net_home(), fname_decoder),
+        )
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "sam"

rembg/sessions/silueta.py

@@ -0,0 +1,49 @@
+import os
+from typing import List
+
+import numpy as np
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class SiluetaSession(BaseSession):
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(
+                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
+            ),
+        )
+
+        pred = ort_outs[0][:, 0, :, :]
+
+        ma = np.max(pred)
+        mi = np.min(pred)
+
+        pred = (pred - mi) / (ma - mi)
+        pred = np.squeeze(pred)
+
+        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+        mask = mask.resize(img.size, Image.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname = f"{cls.name()}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx",
+            "md5:55e59e0d8062d2f5d013f4725ee84782",
+            fname=fname,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "silueta"

rembg/sessions/u2net.py

@@ -0,0 +1,49 @@
+import os
+from typing import List
+
+import numpy as np
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class U2netSession(BaseSession):
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(
+                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
+            ),
+        )
+
+        pred = ort_outs[0][:, 0, :, :]
+
+        ma = np.max(pred)
+        mi = np.min(pred)
+
+        pred = (pred - mi) / (ma - mi)
+        pred = np.squeeze(pred)
+
+        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+        mask = mask.resize(img.size, Image.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname = f"{cls.name()}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx",
+            "md5:60024c5c889badc19c04ad937298a77b",
+            fname=fname,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "u2net"

rembg/sessions/u2net_cloth_seg.py

@@ -0,0 +1,110 @@
+import os
+from typing import List
+
+import numpy as np
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+from scipy.special import log_softmax
+
+from .base import BaseSession
+
+pallete1 = [
+    0,
+    0,
+    0,
+    255,
+    255,
+    255,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+]
+
+pallete2 = [
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    255,
+    255,
+    255,
+    0,
+    0,
+    0,
+]
+
+pallete3 = [
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    0,
+    255,
+    255,
+    255,
+]
+
+
+class Unet2ClothSession(BaseSession):
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(
+                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (768, 768)
+            ),
+        )
+
+        pred = ort_outs
+        pred = log_softmax(pred[0], 1)
+        pred = np.argmax(pred, axis=1, keepdims=True)
+        pred = np.squeeze(pred, 0)
+        pred = np.squeeze(pred, 0)
+
+        mask = Image.fromarray(pred.astype("uint8"), mode="L")
+        mask = mask.resize(img.size, Image.LANCZOS)
+
+        masks = []
+
+        mask1 = mask.copy()
+        mask1.putpalette(pallete1)
+        mask1 = mask1.convert("RGB").convert("L")
+        masks.append(mask1)
+
+        mask2 = mask.copy()
+        mask2.putpalette(pallete2)
+        mask2 = mask2.convert("RGB").convert("L")
+        masks.append(mask2)
+
+        mask3 = mask.copy()
+        mask3.putpalette(pallete3)
+        mask3 = mask3.convert("RGB").convert("L")
+        masks.append(mask3)
+
+        return masks
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname = f"{cls.name()}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx",
+            "md5:2434d1f3cb744e0e49386c906e5a08bb",
+            fname=fname,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "u2net_cloth_seg"

rembg/sessions/u2net_human_seg.py

@@ -0,0 +1,49 @@
+import os
+from typing import List
+
+import numpy as np
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class U2netHumanSegSession(BaseSession):
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(
+                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
+            ),
+        )
+
+        pred = ort_outs[0][:, 0, :, :]
+
+        ma = np.max(pred)
+        mi = np.min(pred)
+
+        pred = (pred - mi) / (ma - mi)
+        pred = np.squeeze(pred)
+
+        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+        mask = mask.resize(img.size, Image.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname = f"{cls.name()}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx",
+            "md5:c09ddc2e0104f800e3e1bb4652583d1f",
+            fname=fname,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "u2net_human_seg"

rembg/sessions/u2netp.py

@@ -0,0 +1,49 @@
+import os
+from typing import List
+
+import numpy as np
+import pooch
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class U2netpSession(BaseSession):
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(
+                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
+            ),
+        )
+
+        pred = ort_outs[0][:, 0, :, :]
+
+        ma = np.max(pred)
+        mi = np.min(pred)
+
+        pred = (pred - mi) / (ma - mi)
+        pred = np.squeeze(pred)
+
+        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+        mask = mask.resize(img.size, Image.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        fname = f"{cls.name()}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx",
+            "md5:8e83ca70e441ab06c318d82300c84806",
+            fname=fname,
+            path=cls.u2net_home(),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        return "u2netp"