add tests

.github/ISSUE_TEMPLATE/bug_report.md

@@ -3,8 +3,7 @@ name: Bug report
 about: Create a report to help us improve
 title: "[BUG] ..."
 labels: bug
-assignees: ''
-
+assignees: ""
 ---
 
 **Describe the bug**
@@ -12,6 +11,7 @@ A clear and concise description of what the bug is.
 
 **To Reproduce**
 Steps to reproduce the behavior:
+
 1. Go to '...'
 2. Click on '....'
 3. Scroll down to '....'

.github/ISSUE_TEMPLATE/feature_request.md

@@ -3,8 +3,7 @@ name: Feature request
 about: Suggest an idea for this project
 title: "[FEATURE] ..."
 labels: enhancement
-assignees: ''
-
+assignees: ""
 ---
 
 **Is your feature request related to a problem? Please describe.**

.github/workflows/close_inactive_issues.yml

@@ -1,23 +1,23 @@
 name: Close inactive issues
 
 on:
-  schedule:
-    - cron: "30 1 * * *"
+    schedule:
+        - cron: "30 1 * * *"
 
 jobs:
-  close-issues:
-    runs-on: ubuntu-latest
-    permissions:
-      issues: write
-      pull-requests: write
-    steps:
-      - uses: actions/stale@v5
-        with:
-          days-before-issue-stale: 30
-          days-before-issue-close: 14
-          stale-issue-label: "stale"
-          stale-issue-message: "This issue is stale because it has been open for 30 days with no activity."
-          close-issue-message: "This issue was closed because it has been inactive for 14 days since being marked as stale."
-          days-before-pr-stale: -1
-          days-before-pr-close: -1
-          repo-token: ${{ secrets.GITHUB_TOKEN }}
+    close-issues:
+        runs-on: ubuntu-latest
+        permissions:
+            issues: write
+            pull-requests: write
+        steps:
+            - uses: actions/stale@v5
+              with:
+                  days-before-issue-stale: 30
+                  days-before-issue-close: 14
+                  stale-issue-label: "stale"
+                  stale-issue-message: "This issue is stale because it has been open for 30 days with no activity."
+                  close-issue-message: "This issue was closed because it has been inactive for 14 days since being marked as stale."
+                  days-before-pr-stale: -1
+                  days-before-pr-close: -1
+                  repo-token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/lint_python.yml

@@ -1,18 +1,18 @@
-name: lint_python
+name: Lint
 
 on: [pull_request, push]
 
 jobs:
-  lint_python:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
-      - run: pip install --upgrade pip wheel
-      - run: pip install bandit black flake8 flake8-bugbear flake8-comprehensions isort safety mypy
-      - run: mypy --install-types --non-interactive --ignore-missing-imports ./rembg
-      - run: bandit --recursive --skip B101,B104,B310,B311,B303 --exclude ./rembg/_version.py ./rembg
-      - run: black --force-exclude rembg/_version.py --check --diff ./rembg
-      - run: flake8 ./rembg --count --ignore=B008,C901,E203,E266,E731,F401,F811,F841,W503 --max-line-length=120 --show-source --statistics --exclude ./rembg/_version.py
-      - run: isort --check-only --profile black ./rembg
-      - run: safety check
+    lint_python:
+        runs-on: ubuntu-latest
+        steps:
+            - uses: actions/checkout@v2
+            - uses: actions/setup-python@v2
+            - run: pip install --upgrade pip wheel
+            - run: pip install bandit black flake8 flake8-bugbear flake8-comprehensions isort safety mypy
+            - run: mypy --install-types --non-interactive --ignore-missing-imports ./rembg
+            - run: bandit --recursive --skip B101,B104,B310,B311,B303 --exclude ./rembg/_version.py ./rembg
+            - run: black --force-exclude rembg/_version.py --check --diff ./rembg
+            - run: flake8 ./rembg --count --ignore=B008,C901,E203,E266,E731,F401,F811,F841,W503 --max-line-length=120 --show-source --statistics --exclude ./rembg/_version.py
+            - run: isort --check-only --profile black ./rembg
+            - run: safety check

.github/workflows/publish_docker.yml

@@ -1,32 +1,28 @@
 name: Publish Docker image
 
 on:
-  push:
-    tags:
-      - "v*.*.*"
+    push:
+        tags:
+            - "v*.*.*"
 
 jobs:
-  push_to_registry:
-    name: Push Docker image to Docker Hub
-    runs-on: ubuntu-latest
-    steps:
-      -
-        name: Checkout
-        uses: actions/checkout@v2
-      -
-        name: Login to Docker Hub
-        uses: docker/login-action@v1
-        with:
-          username: ${{ secrets.DOCKER_HUB_USERNAME }}
-          password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}
-      -
-        name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v1
-      -
-        name: Build and push
-        uses: docker/build-push-action@v2
-        with:
-          context: .
-          file: ./Dockerfile
-          push: true
-          tags: ${{ secrets.DOCKER_HUB_USERNAME }}/rembg:latest
+    push_to_registry:
+        name: Push Docker image to Docker Hub
+        runs-on: ubuntu-latest
+        steps:
+            - name: Checkout
+              uses: actions/checkout@v2
+            - name: Login to Docker Hub
+              uses: docker/login-action@v1
+              with:
+                  username: ${{ secrets.DOCKER_HUB_USERNAME }}
+                  password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}
+            - name: Set up Docker Buildx
+              uses: docker/setup-buildx-action@v1
+            - name: Build and push
+              uses: docker/build-push-action@v2
+              with:
+                  context: .
+                  file: ./Dockerfile
+                  push: true
+                  tags: ${{ secrets.DOCKER_HUB_USERNAME }}/rembg:latest

.github/workflows/publish_pypi.yml

@@ -1,28 +1,28 @@
 name: Publish to Pypi
 
 on:
-  push:
-    tags:
-      - "v*.*.*"
+    push:
+        tags:
+            - "v*.*.*"
 
 jobs:
-  push_to_pypi:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
-        with:
-          python-version: 3.9
+    push_to_pypi:
+        runs-on: ubuntu-latest
+        steps:
+            - uses: actions/checkout@v2
+            - uses: actions/setup-python@v2
+              with:
+                  python-version: 3.9
 
-      - name: "Installs dependencies"
-        run: |
-          python3 -m pip install --upgrade pip
-          python3 -m pip install setuptools wheel twine
+            - name: "Installs dependencies"
+              run: |
+                  python3 -m pip install --upgrade pip
+                  python3 -m pip install setuptools wheel twine
 
-      - name: "Builds and uploads to PyPI"
-        run: |
-          python3 setup.py sdist bdist_wheel
-          python3 -m twine upload dist/*
-        env:
-          TWINE_USERNAME: ${{ secrets.PIPY_USERNAME }}
-          TWINE_PASSWORD: ${{ secrets.PIPY_PASSWORD }}
+            - name: "Builds and uploads to PyPI"
+              run: |
+                  python3 setup.py sdist bdist_wheel
+                  python3 -m twine upload dist/*
+              env:
+                  TWINE_USERNAME: ${{ secrets.PIPY_USERNAME }}
+                  TWINE_PASSWORD: ${{ secrets.PIPY_PASSWORD }}

.github/workflows/tests.yml

@@ -0,0 +1,25 @@
+name: Run tests
+
+on: [push]
+
+jobs:
+    build:
+        runs-on: ubuntu-latest
+        strategy:
+            matrix:
+                python-version: ["3.7", "3.8", "3.9", "3.10"]
+
+        steps:
+            - uses: actions/checkout@v3
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@v4
+              with:
+                  python-version: ${{ matrix.python-version }}
+            - name: Install dependencies
+              run: |
+                  python -m pip install --upgrade pip
+                  pip install pytest
+                  pip install -r requirements.txt
+            - name: Test with pytest
+              run: |
+                  pytest

.gitignore

@@ -13,6 +13,7 @@ __pycache__/
 .envrc
 .python-version
 .idea
+.pytest_cache
 
 # due to using tox and pytest
 .tox

MANIFEST.in

@@ -3,6 +3,8 @@ include LICENSE.txt
 include README.md
 include setup.py
 include pyproject.toml
+include requirements.txt
+include requirements-gpu.txt
 
 include versioneer.py
 include rembg/_version.py

README.md

@@ -37,17 +37,16 @@ Rembg is a tool to remove images background. That is it.
 
 **If this project has helped you, please consider making a [donation](https://www.buymeacoffee.com/danielgatis).**
 
-
 ### Installation
 
-#### **!! This library is for Python 3.9 only !!**
-
 CPU support:
+
 ```bash
 pip install rembg
 ```
 
 GPU support:
+
 ```bash
 pip install rembg[gpu]
 ```
@@ -55,16 +54,19 @@ pip install rembg[gpu]
 ### Usage as a cli
 
 Remove the background from a remote image
+
 ```bash
 curl -s http://input.png | rembg i > output.png
 ```
 
 Remove the background from a local file
+
 ```bash
 rembg i path/to/input.png path/to/output.png
 ```
 
 Remove the background from all images in a folder
+
 ```bash
 rembg p path/to/input path/to/output
 ```
@@ -72,6 +74,7 @@ rembg p path/to/input path/to/output
 ### Usage as a server
 
 Start the server
+
 ```bash
 rembg s
 ```
@@ -83,26 +86,34 @@ http://localhost:5000/docs
 ```
 
 Image with background:
+
 ```
 https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
 ```
 
 Image without background:
+
 ```
 http://localhost:5000/?url=https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
 ```
 
 Also you can send the file as a FormData (multipart/form-data):
+
 ```html
-<form action="http://localhost:5000" method="post" enctype="multipart/form-data">
-   <input type="file" name="file"/>
-   <input type="submit" value="upload"/>
+<form
+    action="http://localhost:5000"
+    method="post"
+    enctype="multipart/form-data"
+>
+    <input type="file" name="file" />
+    <input type="submit" value="upload" />
 </form>
 ```
 
 ### Usage as a library
 
 Input and output as bytes
+
 ```python
 from rembg import remove
 
@@ -117,6 +128,7 @@ with open(input_path, 'rb') as i:
 ```
 
 Input and output as a PIL image
+
 ```python
 from rembg import remove
 from PIL import Image
@@ -130,6 +142,7 @@ output.save(output_path)
 ```
 
 Input and output as a numpy array
+
 ```python
 from rembg import remove
 import cv2
@@ -151,11 +164,13 @@ docker run -p 5000:5000 danielgatis/rembg s
 ```
 
 Image with background:
+
 ```
 https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
 ```
 
 Image without background:
+
 ```
 http://localhost:5000/?url=https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
 ```
@@ -166,10 +181,10 @@ All models are downloaded and saved in the user home folder in the `.u2net` dire
 
 The available models are:
 
-- u2net ([download](https://drive.google.com/uc?id=1tCU5MM1LhRgGou5OpmpjBQbSrYIUoYab), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for general use cases.
-- u2netp ([download](https://drive.google.com/uc?id=1tNuFmLv0TSNDjYIkjEdeH1IWKQdUA4HR), [source](https://github.com/xuebinqin/U-2-Net)): A lightweight version of u2net model.
-- u2net_human_seg ([download](https://drive.google.com/uc?id=1ZfqwVxu-1XWC1xU1GHIP-FM_Knd_AX5j), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for human segmentation.
-- u2net_cloth_seg ([download](https://drive.google.com/uc?id=15rKbQSXQzrKCQurUjZFg8HqzZad8bcyz), [source](https://github.com/levindabhi/cloth-segmentation)): A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body.
+-   u2net ([download](https://drive.google.com/uc?id=1tCU5MM1LhRgGou5OpmpjBQbSrYIUoYab), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for general use cases.
+-   u2netp ([download](https://drive.google.com/uc?id=1tNuFmLv0TSNDjYIkjEdeH1IWKQdUA4HR), [source](https://github.com/xuebinqin/U-2-Net)): A lightweight version of u2net model.
+-   u2net_human_seg ([download](https://drive.google.com/uc?id=1ZfqwVxu-1XWC1xU1GHIP-FM_Knd_AX5j), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for human segmentation.
+-   u2net_cloth_seg ([download](https://drive.google.com/uc?id=15rKbQSXQzrKCQurUjZFg8HqzZad8bcyz), [source](https://github.com/levindabhi/cloth-segmentation)): A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body.
 
 #### How to train your own model
 
@@ -179,6 +194,7 @@ https://github.com/danielgatis/rembg/issues/193#issuecomment-1055534289
 ### Advance usage
 
 Sometimes it is possible to achieve better results by turning on alpha matting. Example:
+
 ```bash
 curl -s http://input.png | rembg i -a -ae 15 > output.png
 ```
@@ -206,11 +222,12 @@ Please contact me at [email protected] if you need help to put it on the clo
 
 ### References
 
-- https://arxiv.org/pdf/2005.09007.pdf
-- https://github.com/NathanUA/U-2-Net
-- https://github.com/pymatting/pymatting
+-   https://arxiv.org/pdf/2005.09007.pdf
+-   https://github.com/NathanUA/U-2-Net
+-   https://github.com/pymatting/pymatting
 
 ### Buy me a coffee
+
 Liked some of my work? Buy me a coffee (or more likely a beer)
 
 <a href="https://www.buymeacoffee.com/danielgatis" target="_blank"><img src="https://bmc-cdn.nyc3.digitaloceanspaces.com/BMC-button-images/custom_images/orange_img.png" alt="Buy Me A Coffee" style="height: auto !important;width: auto !important;"></a>

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+filterwarnings =
+    ignore::DeprecationWarning

rembg/__init__.py

@@ -1,9 +1,3 @@
-import sys
-import warnings
-
-if not (sys.version_info.major == 3 and sys.version_info.minor == 9):
-    warnings.warn("This library is only for Python 3.9", RuntimeWarning)
-
 from . import _version
 
 __version__ = _version.get_versions()["version"]

rembg/bg.py

@@ -16,7 +16,7 @@ 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
 from pymatting.util.util import stack_images
-from scipy.ndimage.morphology import binary_erosion
+from scipy.ndimage import binary_erosion
 
 from .session_base import BaseSession
 from .session_factory import new_session

rembg/session_base.py

@@ -18,7 +18,7 @@ class BaseSession:
         std: Tuple[float, float, float],
         size: Tuple[int, int],
     ) -> 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/session_simple.py

@@ -25,6 +25,6 @@ class SimpleSession(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]

requirements-gpu.txt

@@ -1 +1 @@
-onnxruntime-gpu==1.10.0
+onnxruntime-gpu==1.12.1

requirements.txt

@@ -1,17 +1,18 @@
 aiohttp==3.8.1
 asyncer==0.0.1
-click==8.0.3
-fastapi==0.72.0
-filetype==1.0.9
+click==8.1.3
+fastapi==0.80.0
+filetype==1.1.0
 gdown==4.5.1
-numpy==1.22.3
-onnxruntime==1.12.0
+imagehash==4.2.1
+numpy==1.21.6
+onnxruntime==1.12.1
 opencv-python-headless==4.6.0.66
-pillow==9.0.1
-pymatting==1.1.7
+pillow==9.2.0
+pymatting==1.1.8
 python-multipart==0.0.5
-scikit-image==0.19.1
-scipy==1.8.0
-tqdm==4.62.3
-uvicorn==0.17.0
-watchdog==2.1.7
+scikit-image==0.19.3
+scipy==1.7.3
+tqdm==4.64.0
+uvicorn==0.18.3
+watchdog==2.1.9

setup.py

@@ -11,10 +11,10 @@ here = pathlib.Path(__file__).parent.resolve()
 
 long_description = (here / "README.md").read_text(encoding="utf-8")
 
-with open("requirements.txt") as f:
+with open(here / "requirements.txt") as f:
     requireds = f.read().splitlines()
 
-with open("requirements-gpu.txt") as f:
+with open(here / "requirements-gpu.txt") as f:
     gpu_requireds = f.read().splitlines()
 
 setup(
@@ -27,11 +27,10 @@ setup(
     author_email="[email protected]",
     classifiers=[
         "License :: OSI Approved :: MIT License",
-        "Programming Language :: Python :: 3.9",
     ],
     keywords="remove, background, u2net",
     packages=["rembg"],
-    python_requires="~=3.9.0",
+    python_requires=">=3.7",
     install_requires=requireds,
     entry_points={
         "console_scripts": [

tests/test_remove.py

@@ -0,0 +1,20 @@
+from io import BytesIO
+from pathlib import Path
+
+from imagehash import average_hash
+from PIL import Image
+
+from rembg import remove
+
+here = Path(__file__).parent.resolve()
+
+
+def test_remove():
+    image = Path(here / ".." / "examples" / "animal-1.jpg").read_bytes()
+    expected = Path(here / ".." / "examples" / "animal-1.out.png").read_bytes()
+    actual = remove(image)
+
+    actual_hash = average_hash(Image.open(BytesIO(actual)))
+    expected_hash = average_hash(Image.open(BytesIO(expected)))
+
+    assert actual_hash == expected_hash