initial

.editorconfig

@@ -0,0 +1,11 @@
+# https://editorconfig.org/
+
+root = true
+
+[*]
+indent_style = space
+indent_size = 4
+insert_final_newline = true
+trim_trailing_whitespace = true
+end_of_line = lf
+charset = utf-8

.gitignore

@@ -0,0 +1,13 @@
+# general things to ignore
+build/
+dist/
+*.egg-info/
+*.egg
+*.py[cod]
+__pycache__/
+*.so
+*~
+
+# due to using tox and pytest
+.tox
+.cache

LICENSE.txt

@@ -0,0 +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
+SOFTWARE.

MANIFEST.in

@@ -0,0 +1,12 @@
+include pyproject.toml
+
+# Include the README
+include *.md
+
+# Include the license file
+include LICENSE.txt
+
+# Include the data files
+recursive-include data *
+
+include requirements.txt

README.md

@@ -0,0 +1,96 @@
+
+# Rembg
+
+Rembg is a tool to remove images background. That is it.
+
+<p style="display: flex;align-items: center;justify-content: center;">
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-1.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-1.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-2.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-2.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-3.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-3.out.png" width="100" />
+</p>
+
+<p style="display: flex;align-items: center;justify-content: center;">
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/animal-1.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/animal-1.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/animal-2.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/animal-2.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/animal-3.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/animal-3.out.png" width="100" />
+</p>
+
+<p style="display: flex;align-items: center;justify-content: center;">
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/girl-1.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/girl-1.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/girl-2.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/girl-2.out.png" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/girl-3.jpg" width="100" />
+  <img src="https://raw.githubusercontent.com/danielgatis/rembg/master/examples/girl-3.out.png" width="100" />
+</p>
+
+### Installation
+
+Install it from pypi
+
+```bash
+    pip install rembg
+```
+
+### Usage as a cli
+
+Remove the backaground from a remote image
+```bash
+    curl -s http://input.png | rembg > output.png
+```
+
+Remove the backaground from a local file
+```bash
+    rembg -o path/to/output.png paht/to/input.png
+```
+
+Remove the backaground from all images in a folder
+```bash
+    rembg -p path/to/inputs
+```
+
+### Usage as a server
+
+Start the server
+```bash
+    rembg-server
+```
+
+Open your browser to
+```
+    http://localhost:5000?url=http://image.png
+```
+
+### Usage as a library
+
+In `app.py`
+
+```python
+import sys
+from rembg.bg import remove
+
+sys.stdout.buffer.write(remove(sys.stdin.buffer.read()))
+
+```
+
+Then run
+```
+    cat input.png | python app.py > out.png
+```
+
+### References
+
+- https://arxiv.org/pdf/2005.09007.pdf
+- https://github.com/NathanUA/U-2-Net
+
+### License
+
+Copyright (c) 2020-present [Daniel Gatis](https://github.com/danielgatis)
+
+Licensed under [MIT License](./LICENSE.txt)

pyproject.toml

@@ -0,0 +1,5 @@
+[build-system]
+# These are the assumed default build requirements from pip:
+# https://pip.pypa.io/en/stable/reference/pip/#pep-517-and-518-support
+requires = ["setuptools>=40.8.0", "wheel"]
+build-backend = "setuptools.build_meta"

requirements.txt

@@ -0,0 +1,8 @@
+flask==1.1.2
+numpy==1.19.1
+pillow==7.2.0
+scikit-image==0.17.2
+torch==1.6.0
+torchvision==0.7.0
+waitress==1.4.4
+tqdm==4.48.2

setup.cfg

@@ -0,0 +1,4 @@
+[metadata]
+# This includes the license file(s) in the wheel.
+# https://wheel.readthedocs.io/en/stable/user_guide.html#including-license-files-in-the-generated-wheel-file
+license_files = LICENSE.txt

setup.py

@@ -0,0 +1,36 @@
+import pathlib
+
+from setuptools import find_packages, setup
+
+here = pathlib.Path(__file__).parent.resolve()
+
+long_description = (here / "README.md").read_text(encoding="utf-8")
+
+with open("requirements.txt") as f:
+    requireds = f.read().splitlines()
+
+setup(
+    name="rembg",
+    version="1.0.3",
+    description="Remove image background",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/danielgatis/rembg",
+    author="Daniel Gatis",
+    author_email="[email protected]",
+    classifiers=[
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3 :: Only",
+    ],
+    keywords="remove, background, u2net",
+    package_dir={"": "src"},
+    packages=find_packages(where="src"),
+    python_requires=">=3.5, <4",
+    install_requires=requireds,
+    entry_points={
+        "console_scripts": [
+            "rembg=rembg.cmd.cli:main",
+            "rembg-server=rembg.cmd.server:main",
+        ],
+    },
+)

src/rembg/bg.py

@@ -0,0 +1,30 @@
+import argparse
+import io
+import os
+
+import numpy as np
+from PIL import Image
+
+from .u2net import detect
+
+model_u2net = detect.load_model(model_name="u2net")
+model_u2netp = detect.load_model(model_name="u2netp")
+
+
+def remove(data, model_name="u2net"):
+    model = model_u2net
+
+    if model == "u2netp":
+        model = model_u2netp
+
+    img = Image.open(io.BytesIO(data))
+    roi = detect.predict(model, np.array(img))
+    roi = roi.resize((img.size), resample=Image.LANCZOS)
+
+    empty = Image.new("RGBA", (img.size), 0)
+    out = Image.composite(img, empty, roi.convert("L"))
+
+    bio = io.BytesIO()
+    out.save(bio, "PNG")
+
+    return bio.getbuffer()

src/rembg/cmd/cli.py

@@ -0,0 +1,74 @@
+import argparse
+import glob
+import imghdr
+import io
+import os
+
+import numpy as np
+from PIL import Image
+
+from ..bg import remove
+
+
+def main():
+    ap = argparse.ArgumentParser()
+
+    ap.add_argument(
+        "-m",
+        "--model",
+        default="u2net",
+        type=str,
+        choices=("u2net", "u2netp"),
+        help="The model name.",
+    )
+
+    ap.add_argument(
+        "-p", "--path", nargs="+", help="Path of a file or a folder of files.",
+    )
+
+    ap.add_argument(
+        "-o",
+        "--output",
+        nargs="?",
+        default="-",
+        type=argparse.FileType("wb"),
+        help="Path to the output png image.",
+    )
+
+    ap.add_argument(
+        "input",
+        nargs="?",
+        default="-",
+        type=argparse.FileType("rb"),
+        help="Path to the input image.",
+    )
+
+    args = ap.parse_args()
+
+    r = lambda i: i.buffer.read() if hasattr(i, "buffer") else i.read()
+    w = lambda o, data: o.buffer.write(data) if hasattr(o, "buffer") else o.write(data)
+
+    if args.path:
+        full_paths = [os.path.abspath(path) for path in args.path]
+        files = set()
+
+        for path in full_paths:
+            if os.path.isfile(path):
+                files.add(path)
+            else:
+                full_paths += glob.glob(path + "/*")
+
+        for fi in files:
+            if imghdr.what(fi) is None:
+                continue
+
+            with open(fi, "rb") as input:
+                with open(os.path.splitext(fi)[0] + ".out.png", "wb") as output:
+                    w(output, remove(r(input), args.model))
+
+    else:
+        w(args.output, remove(r(args.input), args.model))
+
+
+if __name__ == "__main__":
+    main()

src/rembg/cmd/server.py

@@ -0,0 +1,51 @@
+import argparse
+from io import BytesIO
+from urllib.parse import unquote_plus
+from urllib.request import urlopen
+
+from flask import Flask, request, send_file
+from waitress import serve
+
+from ..bg import remove
+
+
+def index():
+    model = request.args.get("model", type=str, default="u2net")
+    if model not in ("u2net", "u2netp"):
+        return {"error": "invalid query param 'model'"}, 400
+
+    url = request.args.get("url", type=str)
+    if url is None:
+        return {"error": "missing query param 'url'"}, 400
+
+    try:
+        return send_file(
+            BytesIO(remove(urlopen(unquote_plus(url)).read(), model)),
+            mimetype="image/png",
+        )
+    except Exception as e:
+        app.logger.exception(e.message, exc_info=True)
+        return {"error": "oops, something went wrong!"}, 500
+
+
+def main():
+    ap = argparse.ArgumentParser()
+
+    ap.add_argument(
+        "-a", "--addr", default="0.0.0.0", type=str, help="The IP address to bind to.",
+    )
+
+    ap.add_argument(
+        "-p", "--port", default=5000, type=int, help="The port to bind to.",
+    )
+
+    args = ap.parse_args()
+
+    app = Flask(__name__)
+    app.add_url_rule("/", "index", index)
+
+    serve(app, host=args.addr, port=args.port)
+
+
+if __name__ == "__main__":
+    main()