Add bria-rembg

app.py

@@ -1,109 +1,100 @@
-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(share=True)
+import gradio as gr
+import os
+import cv2
+from rembg import new_session, remove
+from rembg.bg import download_models
+
+def inference(file, mask, model, x, y):
+    session = new_session(model)
+
+    output = remove(
+        file,
+        session=session,
+        **{ "sam_prompt": [{"type": "point", "data": [x, y], "label": 1}] },
+        only_mask=(mask == "Mask only")
+    )
+
+    return output
+
+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)
+
+download_models(tuple())
+
+with gr.Blocks() as app:
+    gr.Markdown(f"# {title}")
+    gr.Markdown(description)
+
+    with gr.Row():
+        inputs = gr.Image(type="numpy", label="Input Image")
+        outputs = gr.Image(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",
+                "bria-rmbg",
+                "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(share=True)

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)"
-    git_full = "e740a9681ea32f5c34adce52aa7cc0b4b85bbb11"
-    git_date = "2024-11-20 09:41:13 -0300"
+    git_refnames = " (HEAD -> main, tag: v2.0.69)"
+    git_full = "df72e3dea3f41e543a13991cb05b8a2659ee95c1"
+    git_date = "2025-12-04 18:05:12 -0300"
     keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
     return keywords
 

rembg/bg.py

@@ -1,4 +1,5 @@
 import io
+import sys
 from enum import Enum
 from typing import Any, List, Optional, Tuple, Union, cast
 
@@ -20,7 +21,7 @@ from pymatting.util.util import stack_images
 from scipy.ndimage import binary_erosion
 
 from .session_factory import new_session
-from .sessions import sessions_class
+from .sessions import sessions, sessions_names
 from .sessions.base import BaseSession
 
 ort.set_default_logger_severity(3)
@@ -175,9 +176,8 @@ def apply_background_color(img: PILImage, color: Tuple[int, int, int, int]) -> P
     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)
+    background = Image.new("RGBA", img.size, tuple(color))
+    colored_image = Image.alpha_composite(background, img)
 
     return colored_image
 
@@ -195,12 +195,25 @@ def fix_image_orientation(img: PILImage) -> PILImage:
     return cast(PILImage, ImageOps.exif_transpose(img))
 
 
-def download_models() -> None:
+def download_models(models: tuple[str, ...]) -> None:
     """
     Download models for image processing.
     """
-    for session in sessions_class:
-        session.download_models()
+    if len(models) == 0:
+        print("No models specified, downloading all models")
+        models = tuple(sessions_names)
+
+    for model in models:
+        session = sessions.get(model)
+        if session is None:
+            print(f"Error: no model found: {model}")
+            sys.exit(1)
+        else:
+            print(f"Downloading model: {model}")
+            try:
+                session.download_models()
+            except Exception as e:
+                print(f"Error downloading model: {e}")
 
 
 def remove(
@@ -215,7 +228,7 @@ def remove(
     bgcolor: Optional[Tuple[int, int, int, int]] = None,
     force_return_bytes: bool = False,
     *args: Optional[Any],
-    **kwargs: Optional[Any]
+    **kwargs: Optional[Any],
 ) -> Union[bytes, PILImage, np.ndarray]:
     """
     Remove the background from an input image.

rembg/commands/b_command.py

@@ -6,7 +6,7 @@ import sys
 from typing import IO
 
 import click
-from PIL.Image import Image as PILImage
+import PIL
 
 from ..bg import remove
 from ..session_factory import new_session
@@ -118,10 +118,11 @@ def b_command(
     Returns:
         None
     """
-    try:
-        kwargs.update(json.loads(extras))
-    except Exception:
-        pass
+    if extras:
+        try:
+            kwargs.update(json.loads(extras))
+        except Exception:
+            raise click.BadParameter("extras must be a valid JSON string")
 
     session = new_session(model, **kwargs)
     bytes_per_img = image_width * image_height * 3
@@ -134,7 +135,7 @@ def b_command(
         if not os.path.isdir(output_dir):
             os.makedirs(output_dir, exist_ok=True)
 
-    def img_to_byte_array(img: PILImage) -> bytes:
+    def img_to_byte_array(img: PIL.Image.Image) -> bytes:
         buff = io.BytesIO()
         img.save(buff, format="PNG")
         return buff.getvalue()
@@ -162,7 +163,7 @@ def b_command(
                 if not img_bytes:
                     break
 
-                img = PILImage.frombytes("RGB", (image_width, image_height), img_bytes)
+                img = PIL.Image.frombytes("RGB", (image_width, image_height), img_bytes)
                 output = remove(img, session=session, **kwargs)
 
                 if output_specifier:

rembg/commands/d_command.py

@@ -5,10 +5,11 @@ from ..bg import download_models
 
 @click.command(  # type: ignore
     name="d",
-    help="download all models",
+    help="download models",
 )
-def d_command(*args, **kwargs) -> None:
+@click.argument("models", nargs=-1)
+def d_command(models: tuple[str, ...]) -> None:
     """
-    Download all models
+    Download models
     """
-    download_models()
+    download_models(models)

rembg/commands/p_command.py

@@ -185,8 +185,7 @@ def p_command(
             print(e)
 
     inputs = list(input.glob("**/*"))
-    if not watch:
-        inputs_tqdm = tqdm(inputs)
+    inputs_tqdm = inputs if watch else tqdm(inputs)
 
     for each_input in inputs_tqdm:
         if not each_input.is_dir():

rembg/commands/s_command.py

@@ -197,12 +197,15 @@ def s_command(port: int, host: str, log_level: str, threads: int) -> None:
             except Exception:
                 pass
 
+        session = sessions.get(commons.model)
+        if session is None:
+            session = new_session(commons.model, **kwargs)
+            sessions[commons.model] = session
+
         return Response(
             remove(
                 content,
-                session=sessions.setdefault(
-                    commons.model, new_session(commons.model, **kwargs)
-                ),
+                session=session,
                 alpha_matting=commons.a,
                 alpha_matting_foreground_threshold=commons.af,
                 alpha_matting_background_threshold=commons.ab,
@@ -306,6 +309,7 @@ def s_command(port: int, host: str, log_level: str, threads: int) -> None:
             ],
             gr.components.Image(type="filepath", label="Output"),
             concurrency_limit=3,
+            analytics_enabled=False,
         )
 
         app = gr.mount_gradio_app(app, interface, path="/")

rembg/session_factory.py

@@ -1,5 +1,5 @@
 import os
-from typing import Type
+from typing import Optional, Type
 
 import onnxruntime as ort
 
@@ -8,9 +8,7 @@ from .sessions.base import BaseSession
 from .sessions.u2net import U2netSession
 
 
-def new_session(
-    model_name: str = "u2net", providers=None, *args, **kwargs
-) -> BaseSession:
+def new_session(model_name: str = "u2net", *args, **kwargs) -> BaseSession:
     """
     Create a new session object based on the specified model name.
 
@@ -21,24 +19,30 @@ def new_session(
 
     Parameters:
         model_name (str): The name of the model.
-        providers: The providers for the session.
         *args: Additional positional arguments.
         **kwargs: Additional keyword arguments.
 
+    Raises:
+        ValueError: If no session class with the given `model_name` is found.
+
     Returns:
         BaseSession: The created session object.
     """
-    session_class: Type[BaseSession] = U2netSession
+    session_class: Optional[Type[BaseSession]] = None
 
     for sc in sessions_class:
         if sc.name() == model_name:
             session_class = sc
             break
 
+    if session_class is None:
+        raise ValueError(f"No session class found for model '{model_name}'")
+
     sess_opts = ort.SessionOptions()
 
     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"])
+        threads = int(os.environ["OMP_NUM_THREADS"])
+        sess_opts.inter_op_num_threads = threads
+        sess_opts.intra_op_num_threads = threads
 
-    return session_class(model_name, sess_opts, providers, *args, **kwargs)
+    return session_class(model_name, sess_opts, *args, **kwargs)

rembg/sessions/__init__.py

@@ -1,88 +1,86 @@
 from __future__ import annotations
 
-from typing import List
+from typing import Dict, List
 
 from .base import BaseSession
 
-sessions_class: List[type[BaseSession]] = []
-sessions_names: List[str] = []
+sessions: Dict[str, type[BaseSession]] = {}
 
 from .birefnet_general import BiRefNetSessionGeneral
 
-sessions_class.append(BiRefNetSessionGeneral)
-sessions_names.append(BiRefNetSessionGeneral.name())
+sessions[BiRefNetSessionGeneral.name()] = BiRefNetSessionGeneral
 
 from .birefnet_general_lite import BiRefNetSessionGeneralLite
 
-sessions_class.append(BiRefNetSessionGeneralLite)
-sessions_names.append(BiRefNetSessionGeneralLite.name())
+sessions[BiRefNetSessionGeneralLite.name()] = BiRefNetSessionGeneralLite
 
 from .birefnet_portrait import BiRefNetSessionPortrait
 
-sessions_class.append(BiRefNetSessionPortrait)
-sessions_names.append(BiRefNetSessionPortrait.name())
+sessions[BiRefNetSessionPortrait.name()] = BiRefNetSessionPortrait
 
 from .birefnet_dis import BiRefNetSessionDIS
 
-sessions_class.append(BiRefNetSessionDIS)
-sessions_names.append(BiRefNetSessionDIS.name())
+sessions[BiRefNetSessionDIS.name()] = BiRefNetSessionDIS
 
 from .birefnet_hrsod import BiRefNetSessionHRSOD
 
-sessions_class.append(BiRefNetSessionHRSOD)
-sessions_names.append(BiRefNetSessionHRSOD.name())
+sessions[BiRefNetSessionHRSOD.name()] = BiRefNetSessionHRSOD
 
 from .birefnet_cod import BiRefNetSessionCOD
 
-sessions_class.append(BiRefNetSessionCOD)
-sessions_names.append(BiRefNetSessionCOD.name())
+sessions[BiRefNetSessionCOD.name()] = BiRefNetSessionCOD
 
 from .birefnet_massive import BiRefNetSessionMassive
 
-sessions_class.append(BiRefNetSessionMassive)
-sessions_names.append(BiRefNetSessionMassive.name())
+sessions[BiRefNetSessionMassive.name()] = BiRefNetSessionMassive
 
 from .dis_anime import DisSession
 
-sessions_class.append(DisSession)
-sessions_names.append(DisSession.name())
+sessions[DisSession.name()] = DisSession
+
+from .dis_custom import DisCustomSession
+
+sessions[DisCustomSession.name()] = DisCustomSession
 
 from .dis_general_use import DisSession as DisSessionGeneralUse
 
-sessions_class.append(DisSessionGeneralUse)
-sessions_names.append(DisSessionGeneralUse.name())
+sessions[DisSessionGeneralUse.name()] = DisSessionGeneralUse
 
 from .sam import SamSession
 
-sessions_class.append(SamSession)
-sessions_names.append(SamSession.name())
+sessions[SamSession.name()] = SamSession
 
 from .silueta import SiluetaSession
 
-sessions_class.append(SiluetaSession)
-sessions_names.append(SiluetaSession.name())
+sessions[SiluetaSession.name()] = SiluetaSession
 
 from .u2net_cloth_seg import Unet2ClothSession
 
-sessions_class.append(Unet2ClothSession)
-sessions_names.append(Unet2ClothSession.name())
+sessions[Unet2ClothSession.name()] = Unet2ClothSession
 
 from .u2net_custom import U2netCustomSession
 
-sessions_class.append(U2netCustomSession)
-sessions_names.append(U2netCustomSession.name())
+sessions[U2netCustomSession.name()] = U2netCustomSession
 
 from .u2net_human_seg import U2netHumanSegSession
 
-sessions_class.append(U2netHumanSegSession)
-sessions_names.append(U2netHumanSegSession.name())
+sessions[U2netHumanSegSession.name()] = U2netHumanSegSession
 
 from .u2net import U2netSession
 
-sessions_class.append(U2netSession)
-sessions_names.append(U2netSession.name())
+sessions[U2netSession.name()] = U2netSession
 
 from .u2netp import U2netpSession
 
-sessions_class.append(U2netpSession)
-sessions_names.append(U2netpSession.name())
+sessions[U2netpSession.name()] = U2netpSession
+
+from .bria_rmbg import BriaRmBgSession
+
+sessions[BriaRmBgSession.name()] = BriaRmBgSession
+
+from .ben_custom import BenCustomSession
+
+sessions[BenCustomSession.name()] = BenCustomSession
+
+sessions_names = list(sessions.keys())
+sessions_class = list(sessions.values())

rembg/sessions/base.py

@@ -10,31 +10,31 @@ 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,
-        sess_opts: ort.SessionOptions,
-        providers=None,
-        *args,
-        **kwargs
-    ):
+    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
         """Initialize an instance of the BaseSession class."""
         self.model_name = model_name
 
-        self.providers = []
-
-        _providers = ort.get_available_providers()
-        if providers:
-            for provider in providers:
-                if provider in _providers:
-                    self.providers.append(provider)
+        if "providers" in kwargs and isinstance(kwargs["providers"], list):
+            providers = kwargs.pop("providers")
         else:
-            self.providers.extend(_providers)
+            device_type = ort.get_device()
+            if (
+                device_type == "GPU"
+                and "CUDAExecutionProvider" in ort.get_available_providers()
+            ):
+                providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            elif (
+                device_type[0:3] == "GPU"
+                and "ROCMExecutionProvider" in ort.get_available_providers()
+            ):
+                providers = ["ROCMExecutionProvider", "CPUExecutionProvider"]
+            else:
+                providers = ["CPUExecutionProvider"]
 
         self.inner_session = ort.InferenceSession(
             str(self.__class__.download_models(*args, **kwargs)),
-            providers=self.providers,
             sess_options=sess_opts,
+            providers=providers,
         )
 
     def normalize(
@@ -49,7 +49,7 @@ class BaseSession:
         im = img.convert("RGB").resize(size, Image.Resampling.LANCZOS)
 
         im_ary = np.array(im)
-        im_ary = im_ary / np.max(im_ary)
+        im_ary = im_ary / max(np.max(im_ary), 1e-6)
 
         tmpImg = np.zeros((im_ary.shape[0], im_ary.shape[1], 3))
         tmpImg[:, :, 0] = (im_ary[:, :, 0] - mean[0]) / std[0]

rembg/sessions/ben_custom.py

@@ -0,0 +1,92 @@
+import os
+from typing import List
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class BenCustomSession(BaseSession):
+    """This is a class representing a custom session for the Ben model."""
+
+    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
+        """
+        Initialize a new BenCustomSession object.
+
+        Parameters:
+            model_name (str): The name of the model.
+            sess_opts: The session options.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        model_path = kwargs.get("model_path")
+        if model_path is None:
+            raise ValueError("model_path is required")
+
+        super().__init__(model_name, sess_opts, *args, **kwargs)
+
+    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.5, 0.5, 0.5), (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.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:
+            raise ValueError("model_path is required")
+
+        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 "ben_custom"

rembg/sessions/bria_rmbg.py

@@ -0,0 +1,88 @@
+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 BriaRmBgSession(BaseSession):
+    """
+    This class represents a Bria-rmbg-2.0 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(
+                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (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.Resampling.LANCZOS)
+
+        return [mask]
+
+    @classmethod
+    def download_models(cls, *args, **kwargs):
+        """
+        Downloads the BRIA-RMBG 2.0 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/bria-rmbg-2.0.onnx",
+            (
+                None
+                if cls.checksum_disabled(*args, **kwargs)
+                else "sha256:5b486f08200f513f460da46dd701db5fbb47d79b4be4b708a19444bcd4e79958"
+            ),
+            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 Bria-rmbg session.
+
+        Parameters:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            str: The name of the session.
+        """
+        return "bria-rmbg"

rembg/sessions/dis_custom.py

@@ -0,0 +1,91 @@
+import os
+from typing import List
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+from PIL.Image import Image as PILImage
+
+from .base import BaseSession
+
+
+class DisCustomSession(BaseSession):
+    """This is a class representing a custom session for the Dis model."""
+
+    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
+        """
+        Initialize a new DisCustomSession object.
+
+        Parameters:
+            model_name (str): The name of the model.
+            sess_opts: The session options.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        model_path = kwargs.get("model_path")
+        if model_path is None:
+            raise ValueError("model_path is required")
+
+        super().__init__(model_name, sess_opts, *args, **kwargs)
+
+    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.5, 0.5, 0.5), (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.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:
+            raise ValueError("model_path is required")
+
+        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 "dis_custom"

rembg/sessions/dis_general_use.py

@@ -24,7 +24,7 @@ class DisSession(BaseSession):
         """
         ort_outs = self.inner_session.run(
             None,
-            self.normalize(img, (0.485, 0.456, 0.406), (1.0, 1.0, 1.0), (1024, 1024)),
+            self.normalize(img, (0.5, 0.5, 0.5), (1.0, 1.0, 1.0), (1024, 1024)),
         )
 
         pred = ort_outs[0][:, 0, :, :]

rembg/sessions/sam.py

@@ -1,6 +1,6 @@
 import os
 from copy import deepcopy
-from typing import Dict, List, Tuple
+from typing import List
 
 import cv2
 import numpy as np
@@ -87,7 +87,6 @@ class SamSession(BaseSession):
         self,
         model_name: str,
         sess_opts: ort.SessionOptions,
-        providers=None,
         *args,
         **kwargs,
     ):
@@ -102,24 +101,13 @@ class SamSession(BaseSession):
         """
         self.model_name = model_name
 
-        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=valid_providers,
             sess_options=sess_opts,
         )
         self.decoder = ort.InferenceSession(
             str(paths[1]),
-            providers=valid_providers,
             sess_options=sess_opts,
         )
 
@@ -142,7 +130,16 @@ class SamSession(BaseSession):
         Returns:
             List[PILImage]: A list of masks generated by the decoder.
         """
-        prompt = kwargs.get("sam_prompt", "{}")
+        prompt = kwargs.get(
+            "sam_prompt",
+            [
+                {
+                    "type": "point",
+                    "label": 1,
+                    "data": [int(img.width / 2), int(img.height / 2)],
+                }
+            ],
+        )
         schema = {
             "type": "array",
             "items": {

rembg/sessions/u2net.py

@@ -41,7 +41,7 @@ class U2netSession(BaseSession):
         pred = (pred - mi) / (ma - mi)
         pred = np.squeeze(pred)
 
-        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+        mask = Image.fromarray((pred.clip(0, 1) * 255).astype("uint8"), mode="L")
         mask = mask.resize(img.size, Image.Resampling.LANCZOS)
 
         return [mask]

rembg/sessions/u2net_custom.py

@@ -13,21 +13,13 @@ 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
-    ):
+    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *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.
 
@@ -38,7 +30,7 @@ class U2netCustomSession(BaseSession):
         if model_path is None:
             raise ValueError("model_path is required")
 
-        super().__init__(model_name, sess_opts, providers, *args, **kwargs)
+        super().__init__(model_name, sess_opts, *args, **kwargs)
 
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
         """
@@ -86,7 +78,7 @@ class U2netCustomSession(BaseSession):
         """
         model_path = kwargs.get("model_path")
         if model_path is None:
-            return
+            raise ValueError("model_path is required")
 
         return os.path.abspath(os.path.expanduser(model_path))
 

requirements.txt

@@ -1,21 +1,22 @@
-pydantic==2.10.6
-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
-aiohttp==3.8.3
-asyncer==0.0.2
-gradio==3.0.20
-jsonschema==4.16.0
+pydantic==2.10.6
+filetype==1.2.0
+pooch==1.6.0
+imagehash==4.3.1
+numpy==1.23.5
+onnxruntime
+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
+aiohttp==3.8.3
+asyncer==0.0.2
+gradio
+jsonschema==4.16.0
+huggingface-hub==0.34.3