mirror of
https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg
synced 2025-07-31 00:11:59 +08:00
add tests
This commit is contained in:
parent
f2efa5d4c5
commit
e7a8a209db
4
.github/ISSUE_TEMPLATE/bug_report.md
vendored
4
.github/ISSUE_TEMPLATE/bug_report.md
vendored
@ -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 '....'
|
||||||
|
3
.github/ISSUE_TEMPLATE/feature_request.md
vendored
3
.github/ISSUE_TEMPLATE/feature_request.md
vendored
@ -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.**
|
||||||
|
36
.github/workflows/close_inactive_issues.yml
vendored
36
.github/workflows/close_inactive_issues.yml
vendored
@ -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 }}
|
||||||
|
28
.github/workflows/lint_python.yml
vendored
28
.github/workflows/lint_python.yml
vendored
@ -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
|
||||||
|
50
.github/workflows/publish_docker.yml
vendored
50
.github/workflows/publish_docker.yml
vendored
@ -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
|
|
||||||
|
42
.github/workflows/publish_pypi.yml
vendored
42
.github/workflows/publish_pypi.yml
vendored
@ -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
25
.github/workflows/tests.yml
vendored
Normal 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
1
.gitignore
vendored
@ -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
|
||||||
|
@ -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
|
||||||
|
43
README.md
43
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).**
|
**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
3
pytest.ini
Normal file
@ -0,0 +1,3 @@
|
|||||||
|
[pytest]
|
||||||
|
filterwarnings =
|
||||||
|
ignore::DeprecationWarning
|
@ -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"]
|
||||||
|
@ -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
|
||||||
|
@ -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)
|
||||||
|
@ -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]
|
||||||
|
@ -1 +1 @@
|
|||||||
onnxruntime-gpu==1.10.0
|
onnxruntime-gpu==1.12.1
|
||||||
|
@ -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
|
||||||
|
7
setup.py
7
setup.py
@ -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
20
tests/test_remove.py
Normal 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
|
Loading…
x
Reference in New Issue
Block a user