Update to latest version + sam support?

.gitattributes

@@ -1,27 +1,27 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zstandard filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE.txt

@@ -1,21 +1,21 @@
-MIT License
-
-Copyright (c) 2020 Daniel Gatis
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+MIT License
+
+Copyright (c) 2020 Daniel Gatis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.

README.md

@@ -1,12 +1,12 @@
----
-title: Rembg
-emoji: 👀
-colorFrom: pink
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.6.0
-app_file: app.py
-pinned: false
----
-
+---
+title: Rembg
+emoji: 👀
+colorFrom: pink
+colorTo: indigo
+sdk: gradio
+sdk_version: 5.6.0
+app_file: app.py
+pinned: false
+---
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,70 +1,109 @@
-## Modified from Akhaliq Hugging Face Demo
-## https://huggingface.co/akhaliq
-
-import gradio as gr
-import os
-import cv2
-
-def inference(file, mask, model):
-    im = cv2.imread(file, cv2.IMREAD_COLOR)
-    cv2.imwrite(os.path.join("input.png"), im)
-
-    from rembg import new_session, remove
-
-    input_path = 'input.png'
-    output_path = 'output.png'
-
-    with open(input_path, 'rb') as i:
-        with open(output_path, 'wb') as o:
-            input = i.read()
-            output = remove(
-                input, 
-                session = new_session(model), 
-                only_mask = (True if mask == "Mask only" else False)
-            ) 
-            
-            
-            
-            o.write(output)
-    return os.path.join("output.png")
-        
-title = "RemBG"
-description = "Gradio demo for RemBG. To use it, simply upload your image and wait. Read more at the link below."
-article = "<p style='text-align: center;'><a href='https://github.com/danielgatis/rembg' target='_blank'>Github Repo</a></p>"
-
-
-gr.Interface(
-    inference, 
-    [
-        gr.inputs.Image(type="filepath", label="Input"),
-        gr.inputs.Radio(
-            [
-                "Default", 
-                "Mask only"
-            ], 
-            type="value",
-            default="Default",
-            label="Choices"
-        ),
-        gr.inputs.Dropdown([
-            "u2net", 
-            "u2netp", 
-            "u2net_human_seg", 
-            "u2net_cloth_seg", 
-            "silueta",
-            "isnet-general-use",
-            "isnet-anime",
-            "sam",
-            ], 
-            type="value",
-            default="isnet-general-use",
-            label="Models"
-        ),
-    ], 
-    gr.outputs.Image(type="filepath", label="Output"),
-    title=title,
-    description=description,
-    article=article,
-    examples=[["lion.png", "Default", "u2net"], ["girl.jpg", "Default", "u2net"], ["anime-girl.jpg", "Default", "isnet-anime"]],
-    enable_queue=True
-    ).launch()
+import gradio as gr
+import os
+import cv2
+from rembg import new_session, remove
+from rembg.sessions import sessions_class
+
+def inference(file, mask, model, x, y):
+    im = cv2.imread(file, cv2.IMREAD_COLOR)
+    input_path = "input.png"
+    output_path = "output.png"
+    cv2.imwrite(input_path, im)
+
+    with open(input_path, 'rb') as i:
+        with open(output_path, 'wb') as o:
+            input = i.read()
+            session = new_session(model)
+
+            output = remove(
+                input,
+                session=session,
+                **{ "sam_prompt": [{"type": "point", "data": [x, y], "label": 1}] },
+                only_mask=(mask == "Mask only")
+            )
+            o.write(output)
+
+    return output_path
+
+title = "RemBG"
+description = "Gradio demo for **[RemBG](https://github.com/danielgatis/rembg)**. To use it, simply upload your image, select a model, click Process, and wait."
+badge = """
+    <div style="position: fixed; left: 50%; text-align: center;">
+        <a href="https://github.com/danielgatis/rembg" target="_blank" style="text-decoration: none;">
+            <img src="https://img.shields.io/badge/RemBG-Github-blue" alt="RemBG Github" />
+        </a>
+    </div>
+"""
+def get_coords(evt: gr.SelectData) -> tuple:
+    return evt.index[0], evt.index[1]
+
+def show_coords(model: str):
+    visible = model == "sam"
+    return gr.update(visible=visible), gr.update(visible=visible), gr.update(visible=visible)
+
+for session in sessions_class:
+    session.download_models()
+
+with gr.Blocks() as app:
+    gr.Markdown(f"# {title}")
+    gr.Markdown(description)
+
+    with gr.Row():
+        inputs = gr.Image(type="filepath", label="Input Image")
+        outputs = gr.Image(type="filepath", label="Output Image")
+        
+    with gr.Row():
+        mask_option = gr.Radio(
+            ["Default", "Mask only"],
+            value="Default",
+            label="Output Type"
+        )
+        model_selector = gr.Dropdown(
+            [
+                "u2net",
+                "u2netp",
+                "u2net_human_seg",
+                "u2net_cloth_seg",
+                "silueta",
+                "isnet-general-use",
+                "isnet-anime",
+                "sam",
+                "birefnet-general",
+                "birefnet-general-lite",
+                "birefnet-portrait",
+                "birefnet-dis",
+                "birefnet-hrsod",
+                "birefnet-cod",
+                "birefnet-massive"
+            ],
+            value="isnet-general-use",
+            label="Model Selection"
+        )
+
+    extra = gr.Markdown("## Click on the image to capture coordinates (for SAM model)", visible=False)
+
+    x = gr.Number(label="Mouse X Coordinate", visible=False)
+    y = gr.Number(label="Mouse Y Coordinate", visible=False)
+
+    model_selector.change(show_coords, inputs=model_selector, outputs=[x, y, extra])
+    inputs.select(get_coords, None, [x, y])
+
+
+    gr.Button("Process Image").click(
+        inference,
+        inputs=[inputs, mask_option, model_selector, x, y],
+        outputs=outputs
+    )
+
+    gr.Examples(
+        examples=[
+            ["lion.png", "Default", "u2net", None, None],
+            ["girl.jpg", "Default", "u2net", None, None],
+            ["anime-girl.jpg", "Default", "isnet-anime", None, None]
+        ],
+        inputs=[inputs, mask_option, model_selector, x, y],
+        outputs=outputs
+    )
+    gr.HTML(badge)
+
+app.launch()

rembg/_version.py

@@ -23,9 +23,9 @@ def get_keywords():
     # setup.py/versioneer.py will grep for the variable names, so they must
     # each be defined on a line of their own. _version.py will just call
     # get_keywords().
-    git_refnames = " (HEAD -> main, tag: v2.0.43)"
-    git_full = "848a38e4cc5cf41522974dea00848596105b1dfa"
-    git_date = "2023-06-02 09:20:57 -0300"
+    git_refnames = " (HEAD -> main)"
+    git_full = "e740a9681ea32f5c34adce52aa7cc0b4b85bbb11"
+    git_date = "2024-11-20 09:41:13 -0300"
     keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
     return keywords
 

rembg/bg.py

@@ -1,8 +1,9 @@
 import io
 from enum import Enum
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union, cast
 
 import numpy as np
+import onnxruntime as ort
 from cv2 import (
     BORDER_DEFAULT,
     MORPH_ELLIPSE,
@@ -22,6 +23,8 @@ from .session_factory import new_session
 from .sessions import sessions_class
 from .sessions.base import BaseSession
 
+ort.set_default_logger_severity(3)
+
 kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
 
 
@@ -38,14 +41,25 @@ def alpha_matting_cutout(
     background_threshold: int,
     erode_structure_size: int,
 ) -> PILImage:
+    """
+    Perform alpha matting on an image using a given mask and threshold values.
+
+    This function takes a PIL image `img` and a PIL image `mask` as input, along with
+    the `foreground_threshold` and `background_threshold` values used to determine
+    foreground and background pixels. The `erode_structure_size` parameter specifies
+    the size of the erosion structure to be applied to the mask.
+
+    The function returns a PIL image representing the cutout of the foreground object
+    from the original image.
+    """
     if img.mode == "RGBA" or img.mode == "CMYK":
         img = img.convert("RGB")
 
-    img = np.asarray(img)
-    mask = np.asarray(mask)
+    img_array = np.asarray(img)
+    mask_array = np.asarray(mask)
 
-    is_foreground = mask > foreground_threshold
-    is_background = mask < background_threshold
+    is_foreground = mask_array > foreground_threshold
+    is_background = mask_array < background_threshold
 
     structure = None
     if erode_structure_size > 0:
@@ -56,11 +70,11 @@ def alpha_matting_cutout(
     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 = np.full(mask_array.shape, dtype=np.uint8, fill_value=128)
     trimap[is_foreground] = 255
     trimap[is_background] = 0
 
-    img_normalized = img / 255.0
+    img_normalized = img_array / 255.0
     trimap_normalized = trimap / 255.0
 
     alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
@@ -74,12 +88,46 @@ def alpha_matting_cutout(
 
 
 def naive_cutout(img: PILImage, mask: PILImage) -> PILImage:
+    """
+    Perform a simple cutout operation on an image using a mask.
+
+    This function takes a PIL image `img` and a PIL image `mask` as input.
+    It uses the mask to create a new image where the pixels from `img` are
+    cut out based on the mask.
+
+    The function returns a PIL image representing the cutout of the original
+    image using the mask.
+    """
     empty = Image.new("RGBA", (img.size), 0)
     cutout = Image.composite(img, empty, mask)
     return cutout
 
 
+def putalpha_cutout(img: PILImage, mask: PILImage) -> PILImage:
+    """
+    Apply the specified mask to the image as an alpha cutout.
+
+    Args:
+        img (PILImage): The image to be modified.
+        mask (PILImage): The mask to be applied.
+
+    Returns:
+        PILImage: The modified image with the alpha cutout applied.
+    """
+    img.putalpha(mask)
+    return img
+
+
 def get_concat_v_multi(imgs: List[PILImage]) -> PILImage:
+    """
+    Concatenate multiple images vertically.
+
+    Args:
+        imgs (List[PILImage]): The list of images to be concatenated.
+
+    Returns:
+        PILImage: The concatenated image.
+    """
     pivot = imgs.pop(0)
     for im in imgs:
         pivot = get_concat_v(pivot, im)
@@ -87,6 +135,16 @@ def get_concat_v_multi(imgs: List[PILImage]) -> PILImage:
 
 
 def get_concat_v(img1: PILImage, img2: PILImage) -> PILImage:
+    """
+    Concatenate two images vertically.
+
+    Args:
+        img1 (PILImage): The first image.
+        img2 (PILImage): The second image to be concatenated below the first image.
+
+    Returns:
+        PILImage: The concatenated image.
+    """
     dst = Image.new("RGBA", (img1.width, img1.height + img2.height))
     dst.paste(img1, (0, 0))
     dst.paste(img2, (0, img1.height))
@@ -102,11 +160,21 @@ def post_process(mask: np.ndarray) -> np.ndarray:
     """
     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
+    mask = np.where(mask < 127, 0, 255).astype(np.uint8)  # type: ignore
     return mask
 
 
 def apply_background_color(img: PILImage, color: Tuple[int, int, int, int]) -> PILImage:
+    """
+    Apply the specified background color to the image.
+
+    Args:
+        img (PILImage): The image to be modified.
+        color (Tuple[int, int, int, int]): The RGBA color to be applied.
+
+    Returns:
+        PILImage: The modified image with the background color applied.
+    """
     r, g, b, a = color
     colored_image = Image.new("RGBA", img.size, (r, g, b, a))
     colored_image.paste(img, mask=img)
@@ -115,10 +183,22 @@ def apply_background_color(img: PILImage, color: Tuple[int, int, int, int]) -> P
 
 
 def fix_image_orientation(img: PILImage) -> PILImage:
-    return ImageOps.exif_transpose(img)
+    """
+    Fix the orientation of the image based on its EXIF data.
+
+    Args:
+        img (PILImage): The image to be fixed.
+
+    Returns:
+        PILImage: The fixed image.
+    """
+    return cast(PILImage, ImageOps.exif_transpose(img))
 
 
 def download_models() -> None:
+    """
+    Download models for image processing.
+    """
     for session in sessions_class:
         session.download_models()
 
@@ -133,20 +213,49 @@ def remove(
     only_mask: bool = False,
     post_process_mask: bool = False,
     bgcolor: Optional[Tuple[int, int, int, int]] = None,
+    force_return_bytes: bool = False,
     *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):
+    """
+    Remove the background from an input image.
+
+    This function takes in various parameters and returns a modified version of the input image with the background removed. The function can handle input data in the form of bytes, a PIL image, or a numpy array. The function first checks the type of the input data and converts it to a PIL image if necessary. It then fixes the orientation of the image and proceeds to perform background removal using the 'u2net' model. The result is a list of binary masks representing the foreground objects in the image. These masks are post-processed and combined to create a final cutout image. If a background color is provided, it is applied to the cutout image. The function returns the resulting cutout image in the format specified by the input 'return_type' parameter or as python bytes if force_return_bytes is true.
+
+    Parameters:
+        data (Union[bytes, PILImage, np.ndarray]): The input image data.
+        alpha_matting (bool, optional): Flag indicating whether to use alpha matting. Defaults to False.
+        alpha_matting_foreground_threshold (int, optional): Foreground threshold for alpha matting. Defaults to 240.
+        alpha_matting_background_threshold (int, optional): Background threshold for alpha matting. Defaults to 10.
+        alpha_matting_erode_size (int, optional): Erosion size for alpha matting. Defaults to 10.
+        session (Optional[BaseSession], optional): A session object for the 'u2net' model. Defaults to None.
+        only_mask (bool, optional): Flag indicating whether to return only the binary masks. Defaults to False.
+        post_process_mask (bool, optional): Flag indicating whether to post-process the masks. Defaults to False.
+        bgcolor (Optional[Tuple[int, int, int, int]], optional): Background color for the cutout image. Defaults to None.
+        force_return_bytes (bool, optional): Flag indicating whether to return the cutout image as bytes. Defaults to False.
+        *args (Optional[Any]): Additional positional arguments.
+        **kwargs (Optional[Any]): Additional keyword arguments.
+
+    Returns:
+        Union[bytes, PILImage, np.ndarray]: The cutout image with the background removed.
+    """
+    if isinstance(data, bytes) or force_return_bytes:
         return_type = ReturnType.BYTES
-        img = Image.open(io.BytesIO(data))
+        img = cast(PILImage, Image.open(io.BytesIO(cast(bytes, data))))
+    elif isinstance(data, PILImage):
+        return_type = ReturnType.PILLOW
+        img = cast(PILImage, data)
     elif isinstance(data, np.ndarray):
         return_type = ReturnType.NDARRAY
-        img = Image.fromarray(data)
+        img = cast(PILImage, Image.fromarray(data))
     else:
-        raise ValueError("Input type {} is not supported.".format(type(data)))
+        raise ValueError(
+            "Input type {} is not supported. Try using force_return_bytes=True to force python bytes output".format(
+                type(data)
+            )
+        )
+
+    putalpha = kwargs.pop("putalpha", False)
 
     # Fix image orientation
     img = fix_image_orientation(img)
@@ -174,10 +283,15 @@ def remove(
                     alpha_matting_erode_size,
                 )
             except ValueError:
-                cutout = naive_cutout(img, mask)
-
+                if putalpha:
+                    cutout = putalpha_cutout(img, mask)
+                else:
+                    cutout = naive_cutout(img, mask)
         else:
-            cutout = naive_cutout(img, mask)
+            if putalpha:
+                cutout = putalpha_cutout(img, mask)
+            else:
+                cutout = naive_cutout(img, mask)
 
         cutouts.append(cutout)
 

rembg/cli.py

@@ -6,9 +6,11 @@ from .commands import command_functions
 
 @click.group()
 @click.version_option(version=_version.get_versions()["version"])
-def main() -> None:
+def _main() -> None:
     pass
 
 
 for command in command_functions:
-    main.add_command(command)
+    _main.add_command(command)
+
+_main()

rembg/commands/__init__.py

@@ -1,13 +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)
+from .b_command import b_command
+from .d_command import d_command
+from .i_command import i_command
+from .p_command import p_command
+from .s_command import s_command
+
+command_functions.append(b_command)
+command_functions.append(d_command)
+command_functions.append(i_command)
+command_functions.append(p_command)
+command_functions.append(s_command)

rembg/commands/b_command.py

@@ -6,14 +6,14 @@ import sys
 from typing import IO
 
 import click
-from PIL import Image
+from PIL.Image import Image as PILImage
 
 from ..bg import remove
 from ..session_factory import new_session
 from ..sessions import sessions_names
 
 
-@click.command(
+@click.command(  # type: ignore
     name="b",
     help="for a byte stream as input",
 )
@@ -74,7 +74,7 @@ from ..sessions import sessions_names
 @click.option(
     "-bgc",
     "--bgcolor",
-    default=None,
+    default=(0, 0, 0, 0),
     type=(int, int, int, int),
     nargs=4,
     help="Background color (R G B A) to replace the removed background with",
@@ -94,7 +94,7 @@ from ..sessions import sessions_names
     "image_height",
     type=int,
 )
-def rs_command(
+def b_command(
     model: str,
     extras: str,
     image_width: int,
@@ -102,12 +102,28 @@ def rs_command(
     output_specifier: str,
     **kwargs
 ) -> None:
+    """
+    Command-line interface for processing images by removing the background using a specified model and generating a mask.
+
+    This CLI command takes several options and arguments to configure the background removal process and save the processed images.
+
+    Parameters:
+        model (str): The name of the model to use for background removal.
+        extras (str): Additional options in JSON format that can be passed to customize the background removal process.
+        image_width (int): The width of the input images in pixels.
+        image_height (int): The height of the input images in pixels.
+        output_specifier (str): A printf-style specifier for the output filenames. If specified, the processed images will be saved to the specified output directory with filenames generated using the specifier.
+        **kwargs: Additional keyword arguments that can be used to customize the background removal process.
+
+    Returns:
+        None
+    """
     try:
         kwargs.update(json.loads(extras))
     except Exception:
         pass
 
-    session = new_session(model)
+    session = new_session(model, **kwargs)
     bytes_per_img = image_width * image_height * 3
 
     if output_specifier:
@@ -118,7 +134,7 @@ def rs_command(
         if not os.path.isdir(output_dir):
             os.makedirs(output_dir, exist_ok=True)
 
-    def img_to_byte_array(img: Image) -> bytes:
+    def img_to_byte_array(img: PILImage) -> bytes:
         buff = io.BytesIO()
         img.save(buff, format="PNG")
         return buff.getvalue()
@@ -146,7 +162,7 @@ def rs_command(
                 if not img_bytes:
                     break
 
-                img = Image.frombytes("RGB", (image_width, image_height), img_bytes)
+                img = PILImage.frombytes("RGB", (image_width, image_height), img_bytes)
                 output = remove(img, session=session, **kwargs)
 
                 if output_specifier:

rembg/commands/d_command.py

@@ -0,0 +1,14 @@
+import click
+
+from ..bg import download_models
+
+
+@click.command(  # type: ignore
+    name="d",
+    help="download all models",
+)
+def d_command(*args, **kwargs) -> None:
+    """
+    Download all models
+    """
+    download_models()

rembg/commands/i_command.py

@@ -9,7 +9,7 @@ from ..session_factory import new_session
 from ..sessions import sessions_names
 
 
-@click.command(
+@click.command(  # type: ignore
     name="i",
     help="for a file as input",
 )
@@ -70,7 +70,7 @@ from ..sessions import sessions_names
 @click.option(
     "-bgc",
     "--bgcolor",
-    default=None,
+    default=(0, 0, 0, 0),
     type=(int, int, int, int),
     nargs=4,
     help="Background color (R G B A) to replace the removed background with",
@@ -85,9 +85,24 @@ from ..sessions import sessions_names
     type=click.File("wb", lazy=True),
 )
 def i_command(model: str, extras: str, input: IO, output: IO, **kwargs) -> None:
+    """
+    Click command line interface function to process an input file based on the provided options.
+
+    This function is the entry point for the CLI program. It reads an input file, applies image processing operations based on the provided options, and writes the output to a file.
+
+    Parameters:
+        model (str): The name of the model to use for image processing.
+        extras (str): Additional options in JSON format.
+        input: The input file to process.
+        output: The output file to write the processed image to.
+        **kwargs: Additional keyword arguments corresponding to the command line options.
+
+    Returns:
+        None
+    """
     try:
         kwargs.update(json.loads(extras))
     except Exception:
         pass
 
-    output.write(remove(input.read(), session=new_session(model), **kwargs))
+    output.write(remove(input.read(), session=new_session(model, **kwargs), **kwargs))

rembg/commands/p_command.py

@@ -14,7 +14,7 @@ from ..session_factory import new_session
 from ..sessions import sessions_names
 
 
-@click.command(
+@click.command(  # type: ignore
     name="p",
     help="for a folder as input",
 )
@@ -80,10 +80,18 @@ from ..sessions import sessions_names
     show_default=True,
     help="watches a folder for changes",
 )
+@click.option(
+    "-d",
+    "--delete_input",
+    default=False,
+    is_flag=True,
+    show_default=True,
+    help="delete input file after processing",
+)
 @click.option(
     "-bgc",
     "--bgcolor",
-    default=None,
+    default=(0, 0, 0, 0),
     type=(int, int, int, int),
     nargs=4,
     help="Background color (R G B A) to replace the removed background with",
@@ -115,14 +123,36 @@ def p_command(
     input: pathlib.Path,
     output: pathlib.Path,
     watch: bool,
+    delete_input: bool,
     **kwargs,
 ) -> None:
+    """
+    Command-line interface (CLI) program for performing background removal on images in a folder.
+
+    This program takes a folder as input and uses a specified model to remove the background from the images in the folder.
+    It provides various options for configuration, such as choosing the model, enabling alpha matting, setting trimap thresholds, erode size, etc.
+    Additional options include outputting only the mask and post-processing the mask.
+    The program can also watch the input folder for changes and automatically process new images.
+    The resulting images with the background removed are saved in the specified output folder.
+
+    Parameters:
+        model (str): The name of the model to use for background removal.
+        extras (str): Additional options in JSON format.
+        input (pathlib.Path): The path to the input folder.
+        output (pathlib.Path): The path to the output folder.
+        watch (bool): Whether to watch the input folder for changes.
+        delete_input (bool): Whether to delete the input file after processing.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        None
+    """
     try:
         kwargs.update(json.loads(extras))
     except Exception:
         pass
 
-    session = new_session(model)
+    session = new_session(model, **kwargs)
 
     def process(each_input: pathlib.Path) -> None:
         try:
@@ -147,33 +177,48 @@ def p_command(
                     print(
                         f"processed: {each_input.absolute()} -> {each_output.absolute()}"
                     )
+
+            if delete_input:
+                each_input.unlink()
+
         except Exception as e:
             print(e)
 
     inputs = list(input.glob("**/*"))
     if not watch:
-        inputs = tqdm(inputs)
+        inputs_tqdm = tqdm(inputs)
 
-    for each_input in inputs:
+    for each_input in inputs_tqdm:
         if not each_input.is_dir():
             process(each_input)
 
     if watch:
+        should_watch = True
         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"]
+                src_path = cast(str, event.src_path)
+                if (
+                    not (
+                        event.is_directory or event.event_type in ["deleted", "closed"]
+                    )
+                    and pathlib.Path(src_path).exists()
                 ):
-                    process(pathlib.Path(event.src_path))
+                    if src_path.endswith("stop.txt"):
+                        nonlocal should_watch
+                        should_watch = False
+                        pathlib.Path(src_path).unlink()
+                        return
+
+                    process(pathlib.Path(src_path))
 
         event_handler = EventHandler()
-        observer.schedule(event_handler, input, recursive=False)
+        observer.schedule(event_handler, str(input), recursive=False)
         observer.start()
 
         try:
-            while True:
+            while should_watch:
                 time.sleep(1)
 
         finally:

rembg/commands/s_command.py

@@ -19,18 +19,26 @@ from ..sessions import sessions_names
 from ..sessions.base import BaseSession
 
 
-@click.command(
+@click.command(  # type: ignore
     name="s",
     help="for a http server",
 )
 @click.option(
     "-p",
     "--port",
-    default=5000,
+    default=7000,
     type=int,
     show_default=True,
     help="port",
 )
+@click.option(
+    "-h",
+    "--host",
+    default="0.0.0.0",
+    type=str,
+    show_default=True,
+    help="host",
+)
 @click.option(
     "-l",
     "--log_level",
@@ -47,7 +55,13 @@ from ..sessions.base import BaseSession
     show_default=True,
     help="number of worker threads",
 )
-def s_command(port: int, log_level: str, threads: int) -> None:
+def s_command(port: int, host: str, log_level: str, threads: int) -> None:
+    """
+    Command-line interface for running the FastAPI web server.
+
+    This function starts the FastAPI web server with the specified port and log level.
+    If the number of worker threads is specified, it sets the thread limiter accordingly.
+    """
     sessions: dict[str, BaseSession] = {}
     tags_metadata = [
         {
@@ -186,7 +200,9 @@ def s_command(port: int, log_level: str, threads: int) -> None:
         return Response(
             remove(
                 content,
-                session=sessions.setdefault(commons.model, new_session(commons.model)),
+                session=sessions.setdefault(
+                    commons.model, new_session(commons.model, **kwargs)
+                ),
                 alpha_matting=commons.a,
                 alpha_matting_foreground_threshold=commons.af,
                 alpha_matting_background_threshold=commons.ab,
@@ -245,12 +261,27 @@ def s_command(port: int, log_level: str, threads: int) -> None:
         return await asyncify(im_without_bg)(file, commons)  # type: ignore
 
     def gr_app(app):
-        def inference(input_path, model):
+        def inference(input_path, model, *args):
             output_path = "output.png"
+            a, af, ab, ae, om, ppm, cmd_args = args
+
+            kwargs = {
+                "alpha_matting": a,
+                "alpha_matting_foreground_threshold": af,
+                "alpha_matting_background_threshold": ab,
+                "alpha_matting_erode_size": ae,
+                "only_mask": om,
+                "post_process_mask": ppm,
+            }
+
+            if cmd_args:
+                kwargs.update(json.loads(cmd_args))
+            kwargs["session"] = new_session(model, **kwargs)
+
             with open(input_path, "rb") as i:
                 with open(output_path, "wb") as o:
                     input = i.read()
-                    output = remove(input, session=new_session(model))
+                    output = remove(input, **kwargs)
                     o.write(output)
             return os.path.join(output_path)
 
@@ -258,28 +289,33 @@ def s_command(port: int, log_level: str, threads: int) -> None:
             inference,
             [
                 gr.components.Image(type="filepath", label="Input"),
-                gr.components.Dropdown(
-                    [
-                        "u2net",
-                        "u2netp",
-                        "u2net_human_seg",
-                        "u2net_cloth_seg",
-                        "silueta",
-                        "isnet-general-use",
-                        "isnet-anime",
-                    ],
-                    value="u2net",
-                    label="Models",
+                gr.components.Dropdown(sessions_names, value="u2net", label="Models"),
+                gr.components.Checkbox(value=True, label="Alpha matting"),
+                gr.components.Slider(
+                    value=240, minimum=0, maximum=255, label="Foreground threshold"
+                ),
+                gr.components.Slider(
+                    value=10, minimum=0, maximum=255, label="Background threshold"
                 ),
+                gr.components.Slider(
+                    value=40, minimum=0, maximum=255, label="Erosion size"
+                ),
+                gr.components.Checkbox(value=False, label="Only mask"),
+                gr.components.Checkbox(value=True, label="Post process mask"),
+                gr.components.Textbox(label="Arguments"),
             ],
             gr.components.Image(type="filepath", label="Output"),
+            concurrency_limit=3,
         )
 
-        interface.queue(concurrency_count=3)
         app = gr.mount_gradio_app(app, interface, path="/")
         return app
 
-    print(f"To access the API documentation, go to http://localhost:{port}/api")
-    print(f"To access the UI, go to http://localhost:{port}")
+    print(
+        f"To access the API documentation, go to http://{'localhost' if host == '0.0.0.0' else host}:{port}/api"
+    )
+    print(
+        f"To access the UI, go to http://{'localhost' if host == '0.0.0.0' else host}:{port}"
+    )
 
-    uvicorn.run(gr_app(app), host="0.0.0.0", port=port, log_level=log_level)
+    uvicorn.run(gr_app(app), host=host, port=port, log_level=log_level)

rembg/session_factory.py

@@ -11,6 +11,23 @@ from .sessions.u2net import U2netSession
 def new_session(
     model_name: str = "u2net", providers=None, *args, **kwargs
 ) -> BaseSession:
+    """
+    Create a new session object based on the specified model name.
+
+    This function searches for the session class based on the model name in the 'sessions_class' list.
+    It then creates an instance of the session class with the provided arguments.
+    The 'sess_opts' object is created using the 'ort.SessionOptions()' constructor.
+    If the 'OMP_NUM_THREADS' environment variable is set, the 'inter_op_num_threads' option of 'sess_opts' is set to its value.
+
+    Parameters:
+        model_name (str): The name of the model.
+        providers: The providers for the session.
+        *args: Additional positional arguments.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        BaseSession: The created session object.
+    """
     session_class: Type[BaseSession] = U2netSession
 
     for sc in sessions_class:
@@ -22,5 +39,6 @@ def new_session(
 
     if "OMP_NUM_THREADS" in os.environ:
         sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
+        sess_opts.intra_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
 
     return session_class(model_name, sess_opts, providers, *args, **kwargs)

rembg/sessions/__init__.py

@@ -1,22 +1,88 @@
-from importlib import import_module
-from inspect import isclass
-from pathlib import Path
-from pkgutil import iter_modules
+from __future__ import annotations
+
+from typing import List
 
 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())
+sessions_class: List[type[BaseSession]] = []
+sessions_names: List[str] = []
+
+from .birefnet_general import BiRefNetSessionGeneral
+
+sessions_class.append(BiRefNetSessionGeneral)
+sessions_names.append(BiRefNetSessionGeneral.name())
+
+from .birefnet_general_lite import BiRefNetSessionGeneralLite
+
+sessions_class.append(BiRefNetSessionGeneralLite)
+sessions_names.append(BiRefNetSessionGeneralLite.name())
+
+from .birefnet_portrait import BiRefNetSessionPortrait
+
+sessions_class.append(BiRefNetSessionPortrait)
+sessions_names.append(BiRefNetSessionPortrait.name())
+
+from .birefnet_dis import BiRefNetSessionDIS
+
+sessions_class.append(BiRefNetSessionDIS)
+sessions_names.append(BiRefNetSessionDIS.name())
+
+from .birefnet_hrsod import BiRefNetSessionHRSOD
+
+sessions_class.append(BiRefNetSessionHRSOD)
+sessions_names.append(BiRefNetSessionHRSOD.name())
+
+from .birefnet_cod import BiRefNetSessionCOD
+
+sessions_class.append(BiRefNetSessionCOD)
+sessions_names.append(BiRefNetSessionCOD.name())
+
+from .birefnet_massive import BiRefNetSessionMassive
+
+sessions_class.append(BiRefNetSessionMassive)
+sessions_names.append(BiRefNetSessionMassive.name())
+
+from .dis_anime import DisSession
+
+sessions_class.append(DisSession)
+sessions_names.append(DisSession.name())
+
+from .dis_general_use import DisSession as DisSessionGeneralUse
+
+sessions_class.append(DisSessionGeneralUse)
+sessions_names.append(DisSessionGeneralUse.name())
+
+from .sam import SamSession
+
+sessions_class.append(SamSession)
+sessions_names.append(SamSession.name())
+
+from .silueta import SiluetaSession
+
+sessions_class.append(SiluetaSession)
+sessions_names.append(SiluetaSession.name())
+
+from .u2net_cloth_seg import Unet2ClothSession
+
+sessions_class.append(Unet2ClothSession)
+sessions_names.append(Unet2ClothSession.name())
+
+from .u2net_custom import U2netCustomSession
+
+sessions_class.append(U2netCustomSession)
+sessions_names.append(U2netCustomSession.name())
+
+from .u2net_human_seg import U2netHumanSegSession
+
+sessions_class.append(U2netHumanSegSession)
+sessions_names.append(U2netHumanSegSession.name())
+
+from .u2net import U2netSession
+
+sessions_class.append(U2netSession)
+sessions_names.append(U2netSession.name())
+
+from .u2netp import U2netpSession
+
+sessions_class.append(U2netpSession)
+sessions_names.append(U2netpSession.name())

rembg/sessions/base.py

@@ -8,6 +8,8 @@ from PIL.Image import Image as PILImage
 
 
 class BaseSession:
+    """This is a base class for managing a session with a machine learning model."""
+
     def __init__(
         self,
         model_name: str,
@@ -16,6 +18,7 @@ class BaseSession:
         *args,
         **kwargs
     ):
+        """Initialize an instance of the BaseSession class."""
         self.model_name = model_name
 
         self.providers = []
@@ -29,7 +32,7 @@ class BaseSession:
             self.providers.extend(_providers)
 
         self.inner_session = ort.InferenceSession(
-            str(self.__class__.download_models()),
+            str(self.__class__.download_models(*args, **kwargs)),
             providers=self.providers,
             sess_options=sess_opts,
         )
@@ -43,7 +46,7 @@ class BaseSession:
         *args,
         **kwargs
     ) -> 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/sessions/birefnet_cod.py

@@ -0,0 +1,52 @@
+import os
+
+import pooch
+
+from . import BiRefNetSessionGeneral
+
+
+class BiRefNetSessionCOD(BiRefNetSessionGeneral):
+    """
+    This class represents a BiRefNet-COD session, which is a subclass of BiRefNetSessionGeneral.
+    """
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BiRefNet-COD model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:f6d0d21ca89d287f17e7afe9f5fd3b45"
+            ),
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Returns the name of the BiRefNet-COD session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "birefnet-cod"

rembg/sessions/birefnet_dis.py

@@ -0,0 +1,52 @@
+import os
+
+import pooch
+
+from . import BiRefNetSessionGeneral
+
+
+class BiRefNetSessionDIS(BiRefNetSessionGeneral):
+    """
+    This class represents a BiRefNet-DIS session, which is a subclass of BiRefNetSessionGeneral.
+    """
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BiRefNet-DIS model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:2d4d44102b446f33a4ebb2e56c051f2b"
+            ),
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Returns the name of the BiRefNet-DIS session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "birefnet-dis"

rembg/sessions/birefnet_general.py

@@ -0,0 +1,91 @@
+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 BiRefNetSessionGeneral(BaseSession):
+    """
+    This class represents a BiRefNet-General session, which is a subclass of BaseSession.
+    """
+
+    def sigmoid(self, mat):
+        return 1 / (1 + np.exp(-mat))
+
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predicts the output masks for the input image using the inner session.
+
+        Parameters:
+            img (PILImage): The input image.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List[PILImage]: The list of output masks.
+        """
+        ort_outs = self.inner_session.run(
+            None,
+            self.normalize(
+                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (1024, 1024)
+            ),
+        )
+
+        pred = self.sigmoid(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.Resampling.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BiRefNet-General model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:7a35a0141cbbc80de11d9c9a28f52697"
+            ),
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Returns the name of the BiRefNet-General session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "birefnet-general"

rembg/sessions/birefnet_general_lite.py

@@ -0,0 +1,52 @@
+import os
+
+import pooch
+
+from . import BiRefNetSessionGeneral
+
+
+class BiRefNetSessionGeneralLite(BiRefNetSessionGeneral):
+    """
+    This class represents a BiRefNet-General-Lite session, which is a subclass of BiRefNetSessionGeneral.
+    """
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BiRefNet-General-Lite model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:4fab47adc4ff364be1713e97b7e66334"
+            ),
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Returns the name of the BiRefNet-General-Lite session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "birefnet-general-lite"

rembg/sessions/birefnet_hrsod.py

@@ -0,0 +1,52 @@
+import os
+
+import pooch
+
+from . import BiRefNetSessionGeneral
+
+
+class BiRefNetSessionHRSOD(BiRefNetSessionGeneral):
+    """
+    This class represents a BiRefNet-HRSOD session, which is a subclass of BiRefNetSessionGeneral.
+    """
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BiRefNet-HRSOD model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:c017ade5de8a50ff0fd74d790d268dda"
+            ),
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Returns the name of the BiRefNet-HRSOD session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "birefnet-hrsod"

rembg/sessions/birefnet_massive.py

@@ -0,0 +1,52 @@
+import os
+
+import pooch
+
+from . import BiRefNetSessionGeneral
+
+
+class BiRefNetSessionMassive(BiRefNetSessionGeneral):
+    """
+    This class represents a BiRefNet-Massive session, which is a subclass of BiRefNetSessionGeneral.
+    """
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BiRefNet-Massive model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:33e726a2136a3d59eb0fdf613e31e3e9"
+            ),
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Returns the name of the BiRefNet-Massive session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "birefnet-massive"

rembg/sessions/birefnet_portrait.py

@@ -0,0 +1,52 @@
+import os
+
+import pooch
+
+from . import BiRefNetSessionGeneral
+
+
+class BiRefNetSessionPortrait(BiRefNetSessionGeneral):
+    """
+    This class represents a BiRefNet-Portrait session, which is a subclass of BiRefNetSessionGeneral.
+    """
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BiRefNet-Portrait model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
+        pooch.retrieve(
+            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:c3a64a6abf20250d090cd055f12a3b67"
+            ),
+            fname=fname,
+            path=cls.u2net_home(*args, **kwargs),
+            progressbar=True,
+        )
+
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Returns the name of the BiRefNet-Portrait session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "birefnet-portrait"

rembg/sessions/dis_anime.py

@@ -10,7 +10,22 @@ from .base import BaseSession
 
 
 class DisSession(BaseSession):
+    """
+    This class represents a session for object detection.
+    """
+
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Use a pre-trained model to predict the object in the given image.
+
+        Parameters:
+            img (PILImage): The input image.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            List[PILImage]: A list of predicted mask images.
+        """
         ort_outs = self.inner_session.run(
             None,
             self.normalize(img, (0.485, 0.456, 0.406), (1.0, 1.0, 1.0), (1024, 1024)),
@@ -25,25 +40,47 @@ class DisSession(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]
 
     @classmethod
     def download_models(cls, *args, **kwargs):
-        fname = f"{cls.name()}.onnx"
+        """
+        Download the pre-trained models.
+
+        Parameters:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            str: The path of the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-anime.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:6f184e756bb3bd901c8849220a83e38e",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:6f184e756bb3bd901c8849220a83e38e"
+            ),
             fname=fname,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
-        return os.path.join(cls.u2net_home(), fname)
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
 
     @classmethod
     def name(cls, *args, **kwargs):
+        """
+        Get the name of the pre-trained model.
+
+        Parameters:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            str: The name of the pre-trained model.
+        """
         return "isnet-anime"

rembg/sessions/dis_general_use.py

@@ -11,6 +11,17 @@ from .base import BaseSession
 
 class DisSession(BaseSession):
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predicts the mask image for the input image.
+
+        This method takes a PILImage object as input and returns a list of PILImage objects as output. It performs several image processing operations to generate the mask image.
+
+        Parameters:
+            img (PILImage): The input image.
+
+        Returns:
+            List[PILImage]: A list of PILImage objects representing the generated mask image.
+        """
         ort_outs = self.inner_session.run(
             None,
             self.normalize(img, (0.485, 0.456, 0.406), (1.0, 1.0, 1.0), (1024, 1024)),
@@ -25,25 +36,51 @@ class DisSession(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]
 
     @classmethod
     def download_models(cls, *args, **kwargs):
-        fname = f"{cls.name()}.onnx"
+        """
+        Downloads the pre-trained model file.
+
+        This class method downloads the pre-trained model file from a specified URL using the pooch library.
+
+        Parameters:
+            args: Additional positional arguments.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:fc16ebd8b0c10d971d3513d564d01e29",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:fc16ebd8b0c10d971d3513d564d01e29"
+            ),
             fname=fname,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
-        return os.path.join(cls.u2net_home(), fname)
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
 
     @classmethod
     def name(cls, *args, **kwargs):
+        """
+        Returns the name of the model.
+
+        This class method returns the name of the model.
+
+        Parameters:
+            args: Additional positional arguments.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the model.
+        """
         return "isnet-general-use"

rembg/sessions/sam.py

@@ -1,9 +1,12 @@
 import os
-from typing import List
+from copy import deepcopy
+from typing import Dict, List, Tuple
 
+import cv2
 import numpy as np
 import onnxruntime as ort
 import pooch
+from jsonschema import validate
 from PIL import Image
 from PIL.Image import Image as PILImage
 
@@ -15,104 +18,213 @@ def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int):
     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:
+def apply_coords(coords: np.ndarray, original_size, target_length):
     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 = deepcopy(coords).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
+def get_input_points(prompt):
+    points = []
+    labels = []
+
+    for mark in prompt:
+        if mark["type"] == "point":
+            points.append(mark["data"])
+            labels.append(mark["label"])
+        elif mark["type"] == "rectangle":
+            points.append([mark["data"][0], mark["data"][1]])
+            points.append([mark["data"][2], mark["data"][3]])
+            labels.append(2)
+            labels.append(3)
 
-    return img.resize((new_w, new_h))
+    points, labels = np.array(points), np.array(labels)
+    return points, labels
 
 
-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
+def transform_masks(masks, original_size, transform_matrix):
+    output_masks = []
+
+    for batch in range(masks.shape[0]):
+        batch_masks = []
+        for mask_id in range(masks.shape[1]):
+            mask = masks[batch, mask_id]
+            mask = cv2.warpAffine(
+                mask,
+                transform_matrix[:2],
+                (original_size[1], original_size[0]),
+                flags=cv2.INTER_LINEAR,
+            )
+            batch_masks.append(mask)
+        output_masks.append(batch_masks)
+
+    return np.array(output_masks)
 
 
 class SamSession(BaseSession):
-    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
+    """
+    This class represents a session for the Sam model.
+
+    Args:
+        model_name (str): The name of the model.
+        sess_opts (ort.SessionOptions): The session options.
+        *args: Variable length argument list.
+        **kwargs: Arbitrary keyword arguments.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        sess_opts: ort.SessionOptions,
+        providers=None,
+        *args,
+        **kwargs,
+    ):
+        """
+        Initialize a new SamSession with the given model name and session options.
+
+        Args:
+            model_name (str): The name of the model.
+            sess_opts (ort.SessionOptions): The session options.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        """
         self.model_name = model_name
-        paths = self.__class__.download_models()
+
+        valid_providers = []
+        available_providers = ort.get_available_providers()
+
+        for provider in providers or []:
+            if provider in available_providers:
+                valid_providers.append(provider)
+        else:
+            valid_providers.extend(available_providers)
+
+        paths = self.__class__.download_models(*args, **kwargs)
         self.encoder = ort.InferenceSession(
             str(paths[0]),
-            providers=ort.get_available_providers(),
+            providers=valid_providers,
             sess_options=sess_opts,
         )
         self.decoder = ort.InferenceSession(
             str(paths[1]),
-            providers=ort.get_available_providers(),
+            providers=valid_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.
+        """
+        Predict masks for an input image.
+
+        This function takes an image as input and performs various preprocessing steps on the image. It then runs the image through an encoder to obtain an image embedding. The function also takes input labels and points as additional arguments. It concatenates the input points and labels with padding and transforms them. It creates an empty mask input and an indicator for no mask. The function then passes the image embedding, point coordinates, point labels, mask input, and has mask input to a decoder. The decoder generates masks based on the input and returns them as a list of images.
+
+        Parameters:
+            img (PILImage): The input image.
+            *args: Additional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List[PILImage]: A list of masks generated by the decoder.
+        """
+        prompt = kwargs.get("sam_prompt", "{}")
+        schema = {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "type": {"type": "string"},
+                    "label": {"type": "integer"},
+                    "data": {
+                        "type": "array",
+                        "items": {"type": "number"},
+                    },
+                },
+            },
+        }
+
+        validate(instance=prompt, schema=schema)
+
+        target_size = 1024
+        input_size = (684, 1024)
+        encoder_input_name = self.encoder.get_inputs()[0].name
+
+        img = img.convert("RGB")
+        cv_image = np.array(img)
+        original_size = cv_image.shape[:2]
+
+        scale_x = input_size[1] / cv_image.shape[1]
+        scale_y = input_size[0] / cv_image.shape[0]
+        scale = min(scale_x, scale_y)
+
+        transform_matrix = np.array(
+            [
+                [scale, 0, 0],
+                [0, scale, 0],
+                [0, 0, 1],
+            ]
+        )
+
+        cv_image = cv2.warpAffine(
+            cv_image,
+            transform_matrix[:2],
+            (input_size[1], input_size[0]),
+            flags=cv2.INTER_LINEAR,
+        )
+
+        ## encoder
+
+        encoder_inputs = {
+            encoder_input_name: cv_image.astype(np.float32),
+        }
+
+        encoder_output = self.encoder.run(None, encoder_inputs)
+        image_embedding = encoder_output[0]
+
+        embedding = {
+            "image_embedding": image_embedding,
+            "original_size": original_size,
+            "transform_matrix": transform_matrix,
+        }
+
+        ## decoder
+
+        input_points, input_labels = get_input_points(prompt)
         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)
+        onnx_coord = apply_coords(onnx_coord, input_size, target_size).astype(
+            np.float32
+        )
+
+        onnx_coord = np.concatenate(
+            [
+                onnx_coord,
+                np.ones((1, onnx_coord.shape[1], 1), dtype=np.float32),
+            ],
+            axis=2,
+        )
+        onnx_coord = np.matmul(onnx_coord, transform_matrix.T)
+        onnx_coord = onnx_coord[:, :, :2].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)
 
@@ -122,48 +234,110 @@ class SamSession(BaseSession):
             "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),
+            "orig_im_size": np.array(input_size, 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])
-        ]
+        masks, _, _ = self.decoder.run(None, decoder_inputs)
+        inv_transform_matrix = np.linalg.inv(transform_matrix)
+        masks = transform_masks(masks, original_size, inv_transform_matrix)
+
+        mask = np.zeros((masks.shape[2], masks.shape[3], 3), dtype=np.uint8)
+        for m in masks[0, :, :, :]:
+            mask[m > 0.0] = [255, 255, 255]
 
-        return masks
+        return [Image.fromarray(mask).convert("L")]
 
     @classmethod
     def download_models(cls, *args, **kwargs):
-        fname_encoder = f"{cls.name()}_encoder.onnx"
-        fname_decoder = f"{cls.name()}_decoder.onnx"
+        """
+        Class method to download ONNX model files.
+
+        This method is responsible for downloading two ONNX model files from specified URLs and saving them locally. The downloaded files are saved with the naming convention 'name_encoder.onnx' and 'name_decoder.onnx', where 'name' is the value returned by the 'name' method.
+
+        Parameters:
+            cls: The class object.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            tuple: A tuple containing the file paths of the downloaded encoder and decoder models.
+        """
+        model_name = kwargs.get("sam_model", "sam_vit_b_01ec64")
+        quant = kwargs.get("sam_quant", False)
+
+        fname_encoder = f"{model_name}.encoder.onnx"
+        fname_decoder = f"{model_name}.decoder.onnx"
+
+        if quant:
+            fname_encoder = f"{model_name}.encoder.quant.onnx"
+            fname_decoder = f"{model_name}.decoder.quant.onnx"
 
         pooch.retrieve(
-            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-encoder-quant.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:13d97c5c79ab13ef86d67cbde5f1b250",
+            f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/{fname_encoder}",
+            None,
             fname=fname_encoder,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
         pooch.retrieve(
-            "https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-decoder-quant.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:fa3d1c36a3187d3de1c8deebf33dd127",
+            f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/{fname_decoder}",
+            None,
             fname=fname_decoder,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
+        if fname_encoder == "sam_vit_h_4b8939.encoder.onnx" and not os.path.exists(
+            os.path.join(
+                cls.u2net_home(*args, **kwargs), "sam_vit_h_4b8939.encoder_data.bin"
+            )
+        ):
+            content = bytearray()
+
+            for i in range(1, 4):
+                pooch.retrieve(
+                    f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/sam_vit_h_4b8939.encoder_data.{i}.bin",
+                    None,
+                    fname=f"sam_vit_h_4b8939.encoder_data.{i}.bin",
+                    path=cls.u2net_home(*args, **kwargs),
+                    progressbar=True,
+                )
+
+                fbin = os.path.join(
+                    cls.u2net_home(*args, **kwargs),
+                    f"sam_vit_h_4b8939.encoder_data.{i}.bin",
+                )
+                content.extend(open(fbin, "rb").read())
+                os.remove(fbin)
+
+            with open(
+                os.path.join(
+                    cls.u2net_home(*args, **kwargs),
+                    "sam_vit_h_4b8939.encoder_data.bin",
+                ),
+                "wb",
+            ) as fp:
+                fp.write(content)
+
         return (
-            os.path.join(cls.u2net_home(), fname_encoder),
-            os.path.join(cls.u2net_home(), fname_decoder),
+            os.path.join(cls.u2net_home(*args, **kwargs), fname_encoder),
+            os.path.join(cls.u2net_home(*args, **kwargs), fname_decoder),
         )
 
     @classmethod
     def name(cls, *args, **kwargs):
+        """
+        Class method to return a string value.
+
+        This method returns the string value 'sam'.
+
+        Parameters:
+            cls: The class object.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            str: The string value 'sam'.
+        """
         return "sam"

rembg/sessions/silueta.py

@@ -10,7 +10,22 @@ from .base import BaseSession
 
 
 class SiluetaSession(BaseSession):
+    """This is a class representing a SiluetaSession object."""
+
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predict the mask of the input image.
+
+        This method takes an image as input, preprocesses it, and performs a prediction to generate a mask. The generated mask is then post-processed and returned as a list of PILImage objects.
+
+        Parameters:
+            img (PILImage): The input image to be processed.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            List[PILImage]: A list of post-processed masks.
+        """
         ort_outs = self.inner_session.run(
             None,
             self.normalize(
@@ -27,25 +42,51 @@ class SiluetaSession(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]
 
     @classmethod
     def download_models(cls, *args, **kwargs):
+        """
+        Download the pre-trained model file.
+
+        This method downloads the pre-trained model file from a specified URL. The file is saved to the U2NET home directory.
+
+        Parameters:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
         fname = f"{cls.name()}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:55e59e0d8062d2f5d013f4725ee84782",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:55e59e0d8062d2f5d013f4725ee84782"
+            ),
             fname=fname,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
-        return os.path.join(cls.u2net_home(), fname)
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
 
     @classmethod
     def name(cls, *args, **kwargs):
+        """
+        Return the name of the model.
+
+        This method returns the name of the Silueta model.
+
+        Parameters:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            str: The name of the model.
+        """
         return "silueta"

rembg/sessions/u2net.py

@@ -10,7 +10,22 @@ from .base import BaseSession
 
 
 class U2netSession(BaseSession):
+    """
+    This class represents a U2net session, which is a subclass of BaseSession.
+    """
+
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predicts the output masks for the input image using the inner session.
+
+        Parameters:
+            img (PILImage): The input image.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List[PILImage]: The list of output masks.
+        """
         ort_outs = self.inner_session.run(
             None,
             self.normalize(
@@ -27,25 +42,47 @@ class U2netSession(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]
 
     @classmethod
     def download_models(cls, *args, **kwargs):
-        fname = f"{cls.name()}.onnx"
+        """
+        Downloads the U2net model file from a specific URL and saves it.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model file.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:60024c5c889badc19c04ad937298a77b",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:60024c5c889badc19c04ad937298a77b"
+            ),
             fname=fname,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
-        return os.path.join(cls.u2net_home(), fname)
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
 
     @classmethod
     def name(cls, *args, **kwargs):
+        """
+        Returns the name of the U2net session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
         return "u2net"

rembg/sessions/u2net_cloth_seg.py

@@ -9,7 +9,7 @@ from scipy.special import log_softmax
 
 from .base import BaseSession
 
-pallete1 = [
+palette1 = [
     0,
     0,
     0,
@@ -24,7 +24,7 @@ pallete1 = [
     0,
 ]
 
-pallete2 = [
+palette2 = [
     0,
     0,
     0,
@@ -39,7 +39,7 @@ pallete2 = [
     0,
 ]
 
-pallete3 = [
+palette3 = [
     0,
     0,
     0,
@@ -57,6 +57,22 @@ pallete3 = [
 
 class Unet2ClothSession(BaseSession):
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predict the cloth category of an image.
+
+        This method takes an image as input and predicts the cloth category of the image.
+        The method uses the inner_session to make predictions using a pre-trained model.
+        The predicted mask is then converted to an image and resized to match the size of the input image.
+        Depending on the cloth category specified in the method arguments, the method applies different color palettes to the mask and appends the resulting images to a list.
+
+        Parameters:
+            img (PILImage): The input image.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List[PILImage]: A list of images representing the predicted masks.
+        """
         ort_outs = self.inner_session.run(
             None,
             self.normalize(
@@ -71,41 +87,59 @@ class Unet2ClothSession(BaseSession):
         pred = np.squeeze(pred, 0)
 
         mask = Image.fromarray(pred.astype("uint8"), mode="L")
-        mask = mask.resize(img.size, Image.LANCZOS)
+        mask = mask.resize(img.size, Image.Resampling.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)
+        cloth_category = kwargs.get("cc") or kwargs.get("cloth_category")
+
+        def upper_cloth():
+            mask1 = mask.copy()
+            mask1.putpalette(palette1)
+            mask1 = mask1.convert("RGB").convert("L")
+            masks.append(mask1)
+
+        def lower_cloth():
+            mask2 = mask.copy()
+            mask2.putpalette(palette2)
+            mask2 = mask2.convert("RGB").convert("L")
+            masks.append(mask2)
+
+        def full_cloth():
+            mask3 = mask.copy()
+            mask3.putpalette(palette3)
+            mask3 = mask3.convert("RGB").convert("L")
+            masks.append(mask3)
+
+        if cloth_category == "upper":
+            upper_cloth()
+        elif cloth_category == "lower":
+            lower_cloth()
+        elif cloth_category == "full":
+            full_cloth()
+        else:
+            upper_cloth()
+            lower_cloth()
+            full_cloth()
 
         return masks
 
     @classmethod
     def download_models(cls, *args, **kwargs):
-        fname = f"{cls.name()}.onnx"
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:2434d1f3cb744e0e49386c906e5a08bb",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:2434d1f3cb744e0e49386c906e5a08bb"
+            ),
             fname=fname,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
-        return os.path.join(cls.u2net_home(), fname)
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
 
     @classmethod
     def name(cls, *args, **kwargs):

rembg/sessions/u2net_custom.py

@@ -0,0 +1,105 @@
+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
+
+
+class U2netCustomSession(BaseSession):
+    """This is a class representing a custom session for the U2net model."""
+
+    def __init__(
+        self,
+        model_name: str,
+        sess_opts: ort.SessionOptions,
+        providers=None,
+        *args,
+        **kwargs
+    ):
+        """
+        Initialize a new U2netCustomSession object.
+
+        Parameters:
+            model_name (str): The name of the model.
+            sess_opts (ort.SessionOptions): The session options.
+            providers: The providers.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Raises:
+            ValueError: If model_path is None.
+        """
+        model_path = kwargs.get("model_path")
+        if model_path is None:
+            raise ValueError("model_path is required")
+
+        super().__init__(model_name, sess_opts, providers, *args, **kwargs)
+
+    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predict the segmentation mask for the input image.
+
+        Parameters:
+            img (PILImage): The input image.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            List[PILImage]: A list of PILImage objects representing the segmentation mask.
+        """
+        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.Resampling.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Download the model files.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The absolute path to the model files.
+        """
+        model_path = kwargs.get("model_path")
+        if model_path is None:
+            return
+
+        return os.path.abspath(os.path.expanduser(model_path))
+
+    @classmethod
+    def name(cls, *args, **kwargs):
+        """
+        Get the name of the model.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the model.
+        """
+        return "u2net_custom"

rembg/sessions/u2net_human_seg.py

@@ -10,7 +10,22 @@ from .base import BaseSession
 
 
 class U2netHumanSegSession(BaseSession):
+    """
+    This class represents a session for performing human segmentation using the U2Net model.
+    """
+
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predicts human segmentation masks for the input image.
+
+        Parameters:
+            img (PILImage): The input image.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            List[PILImage]: A list of predicted masks.
+        """
         ort_outs = self.inner_session.run(
             None,
             self.normalize(
@@ -27,25 +42,47 @@ class U2netHumanSegSession(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]
 
     @classmethod
     def download_models(cls, *args, **kwargs):
-        fname = f"{cls.name()}.onnx"
+        """
+        Downloads the U2Net model weights.
+
+        Parameters:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            str: The path to the downloaded model weights.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:c09ddc2e0104f800e3e1bb4652583d1f",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:c09ddc2e0104f800e3e1bb4652583d1f"
+            ),
             fname=fname,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
-        return os.path.join(cls.u2net_home(), fname)
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
 
     @classmethod
     def name(cls, *args, **kwargs):
+        """
+        Returns the name of the U2Net model.
+
+        Parameters:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            str: The name of the model.
+        """
         return "u2net_human_seg"

rembg/sessions/u2netp.py

@@ -10,7 +10,18 @@ from .base import BaseSession
 
 
 class U2netpSession(BaseSession):
+    """This class represents a session for using the U2netp model."""
+
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
+        """
+        Predicts the mask for the given image using the U2netp model.
+
+        Parameters:
+            img (PILImage): The input image.
+
+        Returns:
+            List[PILImage]: The predicted mask.
+        """
         ort_outs = self.inner_session.run(
             None,
             self.normalize(
@@ -27,25 +38,39 @@ class U2netpSession(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]
 
     @classmethod
     def download_models(cls, *args, **kwargs):
-        fname = f"{cls.name()}.onnx"
+        """
+        Downloads the U2netp model.
+
+        Returns:
+            str: The path to the downloaded model.
+        """
+        fname = f"{cls.name(*args, **kwargs)}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx",
-            None
-            if cls.checksum_disabled(*args, **kwargs)
-            else "md5:8e83ca70e441ab06c318d82300c84806",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "md5:8e83ca70e441ab06c318d82300c84806"
+            ),
             fname=fname,
             path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 
-        return os.path.join(cls.u2net_home(), fname)
+        return os.path.join(cls.u2net_home(*args, **kwargs), fname)
 
     @classmethod
     def name(cls, *args, **kwargs):
+        """
+        Returns the name of the U2netp model.
+
+        Returns:
+            str: The name of the model.
+        """
         return "u2netp"

requirements.txt

@@ -1,18 +1,20 @@
-aiohttp==3.8.1
-asyncer==0.0.2
-click==8.1.3
-fastapi==0.87.0
-filetype==1.2.0
-pooch==1.6.0
-imagehash==4.3.1
-numpy==1.23.5
-onnxruntime==1.13.1
-opencv-python-headless==4.6.0.66
-pillow==9.3.0
-pymatting==1.1.8
-python-multipart==0.0.5
-scikit-image==0.19.3
-scipy==1.9.3
-tqdm==4.64.1
-uvicorn==0.20.0
-watchdog==2.1.9
+filetype==1.2.0
+pooch==1.6.0
+imagehash==4.3.1
+numpy==1.23.5
+onnxruntime==1.13.1
+opencv-python-headless==4.6.0.66
+pillow==9.3.0
+pymatting==1.1.8
+python-multipart==0.0.5
+scikit-image==0.19.3
+scipy==1.9.3
+tqdm==4.64.1
+uvicorn==0.20.0
+watchdog==2.1.9
+click==8.1.3
+fastapi==0.85.0
+aiohttp==3.8.3
+asyncer==0.0.2
+gradio==3.0.20
+jsonschema==4.16.0