Update

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 import os
 import cv2
 
-def inference(file, af, mask, model):
+def inference(file, mask, model):
     im = cv2.imread(file, cv2.IMREAD_COLOR)
     cv2.imwrite(os.path.join("input.png"), im)
 
@@ -20,7 +20,6 @@ def inference(file, af, mask, model):
             output = remove(
                 input, 
                 session = new_session(model), 
-                alpha_matting_erode_size = af, 
                 only_mask = (True if mask == "Mask only" else False)
             ) 
             
@@ -38,7 +37,6 @@ gr.Interface(
     inference, 
     [
         gr.inputs.Image(type="filepath", label="Input"),
-        gr.inputs.Slider(10, 25, default=10, label="Alpha matting erode size"), 
         gr.inputs.Radio(
             [
                 "Default", 
@@ -55,10 +53,11 @@ gr.Interface(
             "u2net_cloth_seg", 
             "silueta",
             "isnet-general-use",
+            "isnet-anime",
             "sam",
             ], 
             type="value",
-            default="u2net",
+            default="isnet-general-use",
             label="Models"
         ),
     ], 
@@ -66,6 +65,6 @@ gr.Interface(
     title=title,
     description=description,
     article=article,
-    examples=[["lion.png", 10, "Default", "u2net"], ["girl.jpg", 10, "Default", "u2net"]],
+    examples=[["lion.png", "Default", "u2net"], ["girl.jpg", "Default", "u2net"], ["anime-girl.jpg", "Default", "isnet-anime"]],
     enable_queue=True
     ).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)"
-    git_full = "e47b2a0ed405a5a30f42bacb142b107f7a4b6536"
-    git_date = "2023-04-26 20:40:21 -0300"
+    git_refnames = " (HEAD -> main, tag: v2.0.43)"
+    git_full = "848a38e4cc5cf41522974dea00848596105b1dfa"
+    git_date = "2023-06-02 09:20:57 -0300"
     keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
     return keywords
 

rembg/bg.py

@@ -11,7 +11,7 @@ from cv2 import (
     getStructuringElement,
     morphologyEx,
 )
-from PIL import Image
+from PIL import Image, ImageOps
 from PIL.Image import Image as PILImage
 from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
 from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
@@ -19,6 +19,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.base import BaseSession
 
 kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
@@ -113,6 +114,15 @@ def apply_background_color(img: PILImage, color: Tuple[int, int, int, int]) -> P
     return colored_image
 
 
+def fix_image_orientation(img: PILImage) -> PILImage:
+    return ImageOps.exif_transpose(img)
+
+
+def download_models() -> None:
+    for session in sessions_class:
+        session.download_models()
+
+
 def remove(
     data: Union[bytes, PILImage, np.ndarray],
     alpha_matting: bool = False,
@@ -138,6 +148,9 @@ def remove(
     else:
         raise ValueError("Input type {} is not supported.".format(type(data)))
 
+    # Fix image orientation
+    img = fix_image_orientation(img)
+
     if session is None:
         session = new_session("u2net", *args, **kwargs)
 

rembg/commands/b_command.py

@@ -0,0 +1,161 @@
+import asyncio
+import io
+import json
+import os
+import sys
+from typing import IO
+
+import click
+from PIL import Image
+
+from ..bg import remove
+from ..session_factory import new_session
+from ..sessions import sessions_names
+
+
+@click.command(
+    name="b",
+    help="for a byte stream as input",
+)
+@click.option(
+    "-m",
+    "--model",
+    default="u2net",
+    type=click.Choice(sessions_names),
+    show_default=True,
+    show_choices=True,
+    help="model name",
+)
+@click.option(
+    "-a",
+    "--alpha-matting",
+    is_flag=True,
+    show_default=True,
+    help="use alpha matting",
+)
+@click.option(
+    "-af",
+    "--alpha-matting-foreground-threshold",
+    default=240,
+    type=int,
+    show_default=True,
+    help="trimap fg threshold",
+)
+@click.option(
+    "-ab",
+    "--alpha-matting-background-threshold",
+    default=10,
+    type=int,
+    show_default=True,
+    help="trimap bg threshold",
+)
+@click.option(
+    "-ae",
+    "--alpha-matting-erode-size",
+    default=10,
+    type=int,
+    show_default=True,
+    help="erode size",
+)
+@click.option(
+    "-om",
+    "--only-mask",
+    is_flag=True,
+    show_default=True,
+    help="output only the mask",
+)
+@click.option(
+    "-ppm",
+    "--post-process-mask",
+    is_flag=True,
+    show_default=True,
+    help="post process the mask",
+)
+@click.option(
+    "-bgc",
+    "--bgcolor",
+    default=None,
+    type=(int, int, int, int),
+    nargs=4,
+    help="Background color (R G B A) to replace the removed background with",
+)
+@click.option("-x", "--extras", type=str)
+@click.option(
+    "-o",
+    "--output_specifier",
+    type=str,
+    help="printf-style specifier for output filenames (e.g. 'output-%d.png'))",
+)
+@click.argument(
+    "image_width",
+    type=int,
+)
+@click.argument(
+    "image_height",
+    type=int,
+)
+def rs_command(
+    model: str,
+    extras: str,
+    image_width: int,
+    image_height: int,
+    output_specifier: str,
+    **kwargs
+) -> None:
+    try:
+        kwargs.update(json.loads(extras))
+    except Exception:
+        pass
+
+    session = new_session(model)
+    bytes_per_img = image_width * image_height * 3
+
+    if output_specifier:
+        output_dir = os.path.dirname(
+            os.path.abspath(os.path.expanduser(output_specifier))
+        )
+
+        if not os.path.isdir(output_dir):
+            os.makedirs(output_dir, exist_ok=True)
+
+    def img_to_byte_array(img: Image) -> bytes:
+        buff = io.BytesIO()
+        img.save(buff, format="PNG")
+        return buff.getvalue()
+
+    async def connect_stdin_stdout():
+        loop = asyncio.get_event_loop()
+        reader = asyncio.StreamReader()
+        protocol = asyncio.StreamReaderProtocol(reader)
+
+        await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+        w_transport, w_protocol = await loop.connect_write_pipe(
+            asyncio.streams.FlowControlMixin, sys.stdout
+        )
+
+        writer = asyncio.StreamWriter(w_transport, w_protocol, reader, loop)
+        return reader, writer
+
+    async def main():
+        reader, writer = await connect_stdin_stdout()
+
+        idx = 0
+        while True:
+            try:
+                img_bytes = await reader.readexactly(bytes_per_img)
+                if not img_bytes:
+                    break
+
+                img = Image.frombytes("RGB", (image_width, image_height), img_bytes)
+                output = remove(img, session=session, **kwargs)
+
+                if output_specifier:
+                    output.save((output_specifier % idx), format="PNG")
+                else:
+                    writer.write(img_to_byte_array(output))
+
+                idx += 1
+            except asyncio.IncompleteReadError:
+                break
+
+    asyncio.run(main())

rembg/commands/s_command.py

@@ -1,8 +1,11 @@
 import json
-from typing import Annotated, Optional, Tuple, cast
+import os
+import webbrowser
+from typing import Optional, Tuple, cast
 
 import aiohttp
 import click
+import gradio as gr
 import uvicorn
 from asyncer import asyncify
 from fastapi import Depends, FastAPI, File, Form, Query
@@ -70,6 +73,7 @@ def s_command(port: int, log_level: str, threads: int) -> None:
             "url": "https://github.com/danielgatis/rembg/blob/main/LICENSE.txt",
         },
         openapi_tags=tags_metadata,
+        docs_url="/api",
     )
 
     app.add_middleware(
@@ -83,10 +87,10 @@ def s_command(port: int, log_level: str, threads: int) -> None:
     class CommonQueryParams:
         def __init__(
             self,
-            model: Annotated[
-                str, Query(regex=r"(" + "|".join(sessions_names) + ")")
-            ] = Query(
+            model: str = Query(
                 description="Model to use when processing image",
+                regex=r"(" + "|".join(sessions_names) + ")",
+                default="u2net",
             ),
             a: bool = Query(default=False, description="Enable Alpha Matting"),
             af: int = Query(
@@ -128,10 +132,10 @@ def s_command(port: int, log_level: str, threads: int) -> None:
     class CommonQueryPostParams:
         def __init__(
             self,
-            model: Annotated[
-                str, Form(regex=r"(" + "|".join(sessions_names) + ")")
-            ] = Form(
+            model: str = Form(
                 description="Model to use when processing image",
+                regex=r"(" + "|".join(sessions_names) + ")",
+                default="u2net",
             ),
             a: bool = Form(default=False, description="Enable Alpha Matting"),
             af: int = Form(
@@ -190,13 +194,18 @@ def s_command(port: int, log_level: str, threads: int) -> None:
                 only_mask=commons.om,
                 post_process_mask=commons.ppm,
                 bgcolor=commons.bgc,
-                **kwargs
+                **kwargs,
             ),
             media_type="image/png",
         )
 
     @app.on_event("startup")
     def startup():
+        try:
+            webbrowser.open(f"http://localhost:{port}")
+        except Exception:
+            pass
+
         if threads is not None:
             from anyio import CapacityLimiter
             from anyio.lowlevel import RunVar
@@ -204,7 +213,7 @@ def s_command(port: int, log_level: str, threads: int) -> None:
             RunVar("_default_thread_limiter").set(CapacityLimiter(threads))
 
     @app.get(
-        path="/",
+        path="/api/remove",
         tags=["Background Removal"],
         summary="Remove from URL",
         description="Removes the background from an image obtained by retrieving an URL.",
@@ -221,7 +230,7 @@ def s_command(port: int, log_level: str, threads: int) -> None:
                 return await asyncify(im_without_bg)(file, commons)
 
     @app.post(
-        path="/",
+        path="/api/remove",
         tags=["Background Removal"],
         summary="Remove from Stream",
         description="Removes the background from an image sent within the request itself.",
@@ -235,4 +244,42 @@ def s_command(port: int, log_level: str, threads: int) -> None:
     ):
         return await asyncify(im_without_bg)(file, commons)  # type: ignore
 
-    uvicorn.run(app, host="0.0.0.0", port=port, log_level=log_level)
+    def gr_app(app):
+        def inference(input_path, model):
+            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))
+                    o.write(output)
+            return os.path.join(output_path)
+
+        interface = gr.Interface(
+            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.Image(type="filepath", label="Output"),
+        )
+
+        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}")
+
+    uvicorn.run(gr_app(app), host="0.0.0.0", port=port, log_level=log_level)

rembg/session_factory.py

@@ -8,7 +8,9 @@ from .sessions.base import BaseSession
 from .sessions.u2net import U2netSession
 
 
-def new_session(model_name: str = "u2net", *args, **kwargs) -> BaseSession:
+def new_session(
+    model_name: str = "u2net", providers=None, *args, **kwargs
+) -> BaseSession:
     session_class: Type[BaseSession] = U2netSession
 
     for sc in sessions_class:
@@ -21,4 +23,4 @@ def new_session(model_name: str = "u2net", *args, **kwargs) -> BaseSession:
     if "OMP_NUM_THREADS" in os.environ:
         sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
 
-    return session_class(model_name, sess_opts, *args, **kwargs)
+    return session_class(model_name, sess_opts, providers, *args, **kwargs)

rembg/sessions/base.py

@@ -8,11 +8,29 @@ from PIL.Image import Image as PILImage
 
 
 class BaseSession:
-    def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
+    def __init__(
+        self,
+        model_name: str,
+        sess_opts: ort.SessionOptions,
+        providers=None,
+        *args,
+        **kwargs
+    ):
         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)
+        else:
+            self.providers.extend(_providers)
+
         self.inner_session = ort.InferenceSession(
             str(self.__class__.download_models()),
-            providers=ort.get_available_providers(),
+            providers=self.providers,
             sess_options=sess_opts,
         )
 
@@ -46,6 +64,10 @@ class BaseSession:
     def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
         raise NotImplementedError
 
+    @classmethod
+    def checksum_disabled(cls, *args, **kwargs):
+        return os.getenv("MODEL_CHECKSUM_DISABLED", None) is not None
+
     @classmethod
     def u2net_home(cls, *args, **kwargs):
         return os.path.expanduser(

rembg/sessions/dis_anime.py

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

rembg/sessions/dis_general_use.py

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

rembg/sessions/sam.py

@@ -141,17 +141,21 @@ class SamSession(BaseSession):
 
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-encoder-quant.onnx",
-            "md5:13d97c5c79ab13ef86d67cbde5f1b250",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:13d97c5c79ab13ef86d67cbde5f1b250",
             fname=fname_encoder,
-            path=cls.u2net_home(),
+            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",
-            "md5:fa3d1c36a3187d3de1c8deebf33dd127",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:fa3d1c36a3187d3de1c8deebf33dd127",
             fname=fname_decoder,
-            path=cls.u2net_home(),
+            path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 

rembg/sessions/silueta.py

@@ -36,9 +36,11 @@ class SiluetaSession(BaseSession):
         fname = f"{cls.name()}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx",
-            "md5:55e59e0d8062d2f5d013f4725ee84782",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:55e59e0d8062d2f5d013f4725ee84782",
             fname=fname,
-            path=cls.u2net_home(),
+            path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 

rembg/sessions/u2net.py

@@ -36,9 +36,11 @@ class U2netSession(BaseSession):
         fname = f"{cls.name()}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx",
-            "md5:60024c5c889badc19c04ad937298a77b",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:60024c5c889badc19c04ad937298a77b",
             fname=fname,
-            path=cls.u2net_home(),
+            path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 

rembg/sessions/u2net_cloth_seg.py

@@ -97,9 +97,11 @@ class Unet2ClothSession(BaseSession):
         fname = f"{cls.name()}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx",
-            "md5:2434d1f3cb744e0e49386c906e5a08bb",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:2434d1f3cb744e0e49386c906e5a08bb",
             fname=fname,
-            path=cls.u2net_home(),
+            path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 

rembg/sessions/u2net_human_seg.py

@@ -36,9 +36,11 @@ class U2netHumanSegSession(BaseSession):
         fname = f"{cls.name()}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx",
-            "md5:c09ddc2e0104f800e3e1bb4652583d1f",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:c09ddc2e0104f800e3e1bb4652583d1f",
             fname=fname,
-            path=cls.u2net_home(),
+            path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )
 

rembg/sessions/u2netp.py

@@ -36,9 +36,11 @@ class U2netpSession(BaseSession):
         fname = f"{cls.name()}.onnx"
         pooch.retrieve(
             "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx",
-            "md5:8e83ca70e441ab06c318d82300c84806",
+            None
+            if cls.checksum_disabled(*args, **kwargs)
+            else "md5:8e83ca70e441ab06c318d82300c84806",
             fname=fname,
-            path=cls.u2net_home(),
+            path=cls.u2net_home(*args, **kwargs),
             progressbar=True,
         )