add tests

This commit is contained in:
Daniel Gatis 2022-08-24 20:10:40 -03:00
parent f2efa5d4c5
commit e7a8a209db
19 changed files with 180 additions and 123 deletions

View File

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

View File

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

View File

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

View File

@ -1,18 +1,18 @@
name: lint_python name: Lint
on: [pull_request, push] on: [pull_request, push]
jobs: jobs:
lint_python: lint_python:
runs-on: ubuntu-latest runs-on: ubuntu-latest
steps: steps:
- uses: actions/checkout@v2 - uses: actions/checkout@v2
- uses: actions/setup-python@v2 - uses: actions/setup-python@v2
- run: pip install --upgrade pip wheel - run: pip install --upgrade pip wheel
- run: pip install bandit black flake8 flake8-bugbear flake8-comprehensions isort safety mypy - run: pip install bandit black flake8 flake8-bugbear flake8-comprehensions isort safety mypy
- run: mypy --install-types --non-interactive --ignore-missing-imports ./rembg - 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: bandit --recursive --skip B101,B104,B310,B311,B303 --exclude ./rembg/_version.py ./rembg
- run: black --force-exclude rembg/_version.py --check --diff ./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: 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: isort --check-only --profile black ./rembg
- run: safety check - run: safety check

View File

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

View File

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

25
.github/workflows/tests.yml vendored Normal file
View File

@ -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

1
.gitignore vendored
View File

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

View File

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

View File

@ -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).** **If this project has helped you, please consider making a [donation](https://www.buymeacoffee.com/danielgatis).**
### Installation ### Installation
#### **!! This library is for Python 3.9 only !!**
CPU support: CPU support:
```bash ```bash
pip install rembg pip install rembg
``` ```
GPU support: GPU support:
```bash ```bash
pip install rembg[gpu] pip install rembg[gpu]
``` ```
@ -55,16 +54,19 @@ pip install rembg[gpu]
### Usage as a cli ### Usage as a cli
Remove the background from a remote image Remove the background from a remote image
```bash ```bash
curl -s http://input.png | rembg i > output.png curl -s http://input.png | rembg i > output.png
``` ```
Remove the background from a local file Remove the background from a local file
```bash ```bash
rembg i path/to/input.png path/to/output.png rembg i path/to/input.png path/to/output.png
``` ```
Remove the background from all images in a folder Remove the background from all images in a folder
```bash ```bash
rembg p path/to/input path/to/output rembg p path/to/input path/to/output
``` ```
@ -72,6 +74,7 @@ rembg p path/to/input path/to/output
### Usage as a server ### Usage as a server
Start the server Start the server
```bash ```bash
rembg s rembg s
``` ```
@ -83,26 +86,34 @@ http://localhost:5000/docs
``` ```
Image with background: Image with background:
``` ```
https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
``` ```
Image without background: 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 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): Also you can send the file as a FormData (multipart/form-data):
```html ```html
<form action="http://localhost:5000" method="post" enctype="multipart/form-data"> <form
<input type="file" name="file"/> action="http://localhost:5000"
<input type="submit" value="upload"/> method="post"
enctype="multipart/form-data"
>
<input type="file" name="file" />
<input type="submit" value="upload" />
</form> </form>
``` ```
### Usage as a library ### Usage as a library
Input and output as bytes Input and output as bytes
```python ```python
from rembg import remove from rembg import remove
@ -117,6 +128,7 @@ with open(input_path, 'rb') as i:
``` ```
Input and output as a PIL image Input and output as a PIL image
```python ```python
from rembg import remove from rembg import remove
from PIL import Image from PIL import Image
@ -130,6 +142,7 @@ output.save(output_path)
``` ```
Input and output as a numpy array Input and output as a numpy array
```python ```python
from rembg import remove from rembg import remove
import cv2 import cv2
@ -151,11 +164,13 @@ docker run -p 5000:5000 danielgatis/rembg s
``` ```
Image with background: Image with background:
``` ```
https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg https://upload.wikimedia.org/wikipedia/commons/thumb/9/9a/Gull_portrait_ca_usa.jpg/1280px-Gull_portrait_ca_usa.jpg
``` ```
Image without background: 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 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: 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. - 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. - 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_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_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 #### How to train your own model
@ -179,6 +194,7 @@ https://github.com/danielgatis/rembg/issues/193#issuecomment-1055534289
### Advance usage ### Advance usage
Sometimes it is possible to achieve better results by turning on alpha matting. Example: Sometimes it is possible to achieve better results by turning on alpha matting. Example:
```bash ```bash
curl -s http://input.png | rembg i -a -ae 15 > output.png curl -s http://input.png | rembg i -a -ae 15 > output.png
``` ```
@ -206,11 +222,12 @@ Please contact me at danielgatis@gmail.com if you need help to put it on the clo
### References ### References
- https://arxiv.org/pdf/2005.09007.pdf - https://arxiv.org/pdf/2005.09007.pdf
- https://github.com/NathanUA/U-2-Net - https://github.com/NathanUA/U-2-Net
- https://github.com/pymatting/pymatting - https://github.com/pymatting/pymatting
### Buy me a coffee ### Buy me a coffee
Liked some of my work? Buy me a coffee (or more likely a beer) 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> <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>

3
pytest.ini Normal file
View File

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

View File

@ -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 from . import _version
__version__ = _version.get_versions()["version"] __version__ = _version.get_versions()["version"]

View File

@ -16,7 +16,7 @@ from PIL.Image import Image as PILImage
from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
from pymatting.util.util import stack_images 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_base import BaseSession
from .session_factory import new_session from .session_factory import new_session

View File

@ -18,7 +18,7 @@ class BaseSession:
std: Tuple[float, float, float], std: Tuple[float, float, float],
size: Tuple[int, int], size: Tuple[int, int],
) -> Dict[str, np.ndarray]: ) -> 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 = np.array(im)
im_ary = im_ary / np.max(im_ary) im_ary = im_ary / np.max(im_ary)

View File

@ -25,6 +25,6 @@ class SimpleSession(BaseSession):
pred = np.squeeze(pred) pred = np.squeeze(pred)
mask = Image.fromarray((pred * 255).astype("uint8"), mode="L") 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] return [mask]

View File

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

View File

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

View File

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

20
tests/test_remove.py Normal file
View File

@ -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