Repository: danielgatis/rembg
Branch: main
Commit: 5e10e9baaaef
Files: 64
Total size: 141.4 KB
Directory structure:
gitextract_xxyh1rmd/
├── .dockerignore
├── .editorconfig
├── .gitattributes
├── .github/
│ ├── FUNDING.yml
│ ├── ISSUE_TEMPLATE/
│ │ ├── bug_report.md
│ │ └── feature_request.md
│ └── workflows/
│ ├── close_inactive_issues.yml
│ ├── lint_python.yml
│ ├── publish_docker.yml
│ ├── publish_pypi.yml
│ └── windows_installer.yml
├── .gitignore
├── .markdownlint.yaml
├── .python-version
├── CITATION.cff
├── Dockerfile
├── Dockerfile_nvidia_cuda_cudnn_gpu
├── LICENSE.txt
├── MANIFEST.in
├── README.md
├── USAGE.md
├── _build-exe.ps1
├── _modpath.iss
├── _setup-cpu.iss
├── _setup-gpu.iss
├── docker-compose.yml
├── man/
│ └── rembg.1
├── pyproject.toml
├── pytest.ini
├── rembg/
│ ├── __init__.py
│ ├── bg.py
│ ├── cli.py
│ ├── commands/
│ │ ├── __init__.py
│ │ ├── b_command.py
│ │ ├── d_command.py
│ │ ├── i_command.py
│ │ ├── p_command.py
│ │ └── s_command.py
│ ├── session_factory.py
│ └── sessions/
│ ├── __init__.py
│ ├── base.py
│ ├── ben_custom.py
│ ├── birefnet_cod.py
│ ├── birefnet_dis.py
│ ├── birefnet_general.py
│ ├── birefnet_general_lite.py
│ ├── birefnet_hrsod.py
│ ├── birefnet_massive.py
│ ├── birefnet_portrait.py
│ ├── bria_rmbg.py
│ ├── dis_anime.py
│ ├── dis_custom.py
│ ├── dis_general_use.py
│ ├── sam.py
│ ├── silueta.py
│ ├── u2net.py
│ ├── u2net_cloth_seg.py
│ ├── u2net_custom.py
│ ├── u2net_human_seg.py
│ └── u2netp.py
├── rembg.ipynb
├── rembg.py
├── rembg.spec
└── tests/
└── test_remove.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .dockerignore
================================================
*
!rembg
!pyproject.toml
!poetry.lock
!README.md
!.git
.env
================================================
FILE: .editorconfig
================================================
# 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
================================================
FILE: .gitattributes
================================================
rembg/_version.py export-subst
================================================
FILE: .github/FUNDING.yml
================================================
github: [danielgatis]
custom: ["https://www.buymeacoffee.com/danielgatis"]
================================================
FILE: .github/ISSUE_TEMPLATE/bug_report.md
================================================
---
name: Bug report
about: Create a report to help us improve
title: "[BUG] ..."
labels: bug
assignees: ""
---
**Describe the bug**
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 '....'
4. See error
**Expected behavior**
A clear and concise description of what you expected to happen.
**Images**
Input images to reproduce.
**OS Version:**
iOS 22
**Rembg version:**
v2.0.21
**Additional context**
Add any other context about the problem here.
================================================
FILE: .github/ISSUE_TEMPLATE/feature_request.md
================================================
---
name: Feature request
about: Suggest an idea for this project
title: "[FEATURE] ..."
labels: enhancement
assignees: ""
---
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
**Describe the solution you'd like**
A clear and concise description of what you want to happen.
**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.
**Additional context**
Add any other context or screenshots about the feature request here.
================================================
FILE: .github/workflows/close_inactive_issues.yml
================================================
name: Close inactive issues
on:
schedule:
- cron: "30 1 * * *"
jobs:
close_inactive_issues:
runs-on: ubuntu-latest
permissions:
issues: write
pull-requests: write
steps:
- uses: actions/stale@v9
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 }}
================================================
FILE: .github/workflows/lint_python.yml
================================================
name: Lint
on:
push:
branches:
- "**"
pull_request:
jobs:
lint_python:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
- name: Install Poetry
uses: snok/install-poetry@v1
with:
virtualenvs-create: true
virtualenvs-in-project: true
- name: Install dependencies
run: poetry install --with dev --extras "cpu cli"
- run: poetry run mypy --install-types --non-interactive --ignore-missing-imports ./rembg
- run: poetry run bandit --recursive --skip B101,B104,B310,B311,B303,B110 --exclude ./rembg/_version.py ./rembg
- run: poetry run black --force-exclude rembg/_version.py --check --diff ./rembg
- run: poetry run flake8 ./rembg --count --ignore=B008,C901,E203,E266,E731,F401,F811,F841,W503,E501,E402 --show-source --statistics --exclude ./rembg/_version.py
- run: poetry run isort --check-only --profile black ./rembg
================================================
FILE: .github/workflows/publish_docker.yml
================================================
name: Publish Docker image
on:
push:
tags:
- "v*.*.*"
jobs:
publish_docker:
name: Push Docker image to Docker Hub
runs-on: ubuntu-24.04
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Docker meta
id: meta
uses: docker/metadata-action@v5
with:
# list of Docker images to use as base name for tags
images: |
${{ secrets.DOCKER_HUB_USERNAME }}/rembg
# generate Docker tags based on the following events/attributes
tags: |
type=ref,event=branch
type=ref,event=branch
type=ref,event=pr
type=semver,pattern={{version}}
type=semver,pattern={{major}}.{{minor}}
type=semver,pattern={{major}}
type=sha
- name: Set up QEMU
uses: docker/setup-qemu-action@v3
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- name: Login to Docker Hub
uses: docker/login-action@v3
with:
username: ${{ secrets.DOCKER_HUB_USERNAME }}
password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}
- name: Build and push
uses: docker/build-push-action@v6
with:
context: .
platforms: linux/amd64
push: ${{ github.event_name != 'pull_request' }}
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=registry,ref=${{ secrets.DOCKER_HUB_USERNAME }}/rembg:buildcache
cache-to: type=registry,ref=${{ secrets.DOCKER_HUB_USERNAME }}/rembg:buildcache,mode=max
================================================
FILE: .github/workflows/publish_pypi.yml
================================================
name: Publish to Pypi
on:
push:
tags:
- "v*.*.*"
jobs:
publish_pypi:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 0
- uses: actions/setup-python@v5
- name: Install Poetry
uses: snok/install-poetry@v1
with:
virtualenvs-create: true
virtualenvs-in-project: true
- name: Install dependencies
run: |
poetry self add "poetry-dynamic-versioning[plugin]"
poetry install --with dev --extras "cpu cli"
- name: Build and publish to PyPI
run: |
poetry build
poetry publish
env:
POETRY_PYPI_TOKEN_PYPI: ${{ secrets.PIPY_PASSWORD }}
test_install:
needs: publish_pypi
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.13"]
steps:
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: Wait for PyPI to update
run: sleep 60
- name: Install from PyPI
run: pip install rembg[cpu,cli]
- name: Test installation
run: |
attempt=0
until rembg d || [ $attempt -eq 5 ]; do
attempt=$((attempt+1))
echo "Attempt $attempt to download the models..."
done
if [ $attempt -eq 5 ]; then
echo "downloading the models failed 5 times, exiting..."
exit 1
fi
rembg --version
================================================
FILE: .github/workflows/windows_installer.yml
================================================
name: Build Windows Installer
on:
push:
tags:
- "v*.*.*"
jobs:
windows_installer:
name: Build the Inno Setup Installer
runs-on: windows-latest
steps:
- uses: actions/setup-python@v5
- uses: actions/checkout@v4
- shell: pwsh
run: ./_build-exe.ps1
- name: Compile CPU Installer
uses: Minionguyjpro/Inno-Setup-Action@v1.2.2
with:
path: _setup-cpu.iss
options: /O+
- name: Compile GPU Installer
uses: Minionguyjpro/Inno-Setup-Action@v1.2.2
with:
path: _setup-gpu.iss
options: /O+
- name: Upload CPU installer to release
uses: svenstaro/upload-release-action@v2
with:
repo_token: ${{ secrets.GITHUB_TOKEN }}
file: dist/rembg-cli-cpu-installer.exe
asset_name: rembg-cli-cpu-installer.exe
tag: ${{ github.ref }}
overwrite: true
- name: Upload GPU installer to release
uses: svenstaro/upload-release-action@v2
with:
repo_token: ${{ secrets.GITHUB_TOKEN }}
file: dist/rembg-cli-gpu-installer.exe
asset_name: rembg-cli-gpu-installer.exe
tag: ${{ github.ref }}
overwrite: true
================================================
FILE: .gitignore
================================================
# general things to ignore
build/
dist/
.venv/
.direnv/
*.egg-info/
*.egg
*.py[cod]
__pycache__/
*.so
*~≈
.env
.envrc
.idea
.pytest_cache
# due to using tox and pytest
.tox
.cache
.mypy_cache
# Poetry
# For libraries, poetry.lock is often not committed
# For applications, it should be committed
poetry.lock
================================================
FILE: .markdownlint.yaml
================================================
---
default: true
MD013: false # line-length
MD033: false # no-inline-html
================================================
FILE: .python-version
================================================
3.13.9
================================================
FILE: CITATION.cff
================================================
cff-version: 1.2.0
title: rembg
message: Rembg is a tool to remove images background
type: software
authors:
- given-names: Daniel
family-names: Gatis
email: danielgatis@gmail.com
identifiers:
- type: url
value: 'https://github.com/danielgatis'
repository-code: 'https://github.com/danielgatis/rembg'
url: 'https://github.com/danielgatis/rembg'
abstract: Rembg is a tool to remove images background.
license: MIT
commit: 9079508935ae55d6eefa0fd75f870599640e8593
version: 2.0.66
date-released: '2025-02-21'
================================================
FILE: Dockerfile
================================================
FROM python:3.11-slim
WORKDIR /rembg
RUN pip install --upgrade pip && \
pip install poetry poetry-dynamic-versioning
RUN apt-get update && apt-get install -y curl git && apt-get clean && rm -rf /var/lib/apt/lists/*
COPY . .
RUN poetry config virtualenvs.create false && \
poetry install --extras "cpu cli" --without dev
RUN rembg d u2net
EXPOSE 7000
ENTRYPOINT ["rembg"]
CMD ["--help"]
================================================
FILE: Dockerfile_nvidia_cuda_cudnn_gpu
================================================
FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
WORKDIR /rembg
RUN apt-get update && apt-get install -y --no-install-recommends python3-pip python-is-python3 curl && apt-get clean && rm -rf /var/lib/apt/lists/*
COPY . .
RUN python -m pip install ".[gpu,cli]" --break-system-packages
RUN rembg d u2net
EXPOSE 7000
ENTRYPOINT ["rembg"]
CMD ["--help"]
================================================
FILE: LICENSE.txt
================================================
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.
================================================
FILE: MANIFEST.in
================================================
include LICENSE.txt
include README.md
include pyproject.toml
================================================
FILE: README.md
================================================
<p align="center">
<img src="logo.png" alt="Rembg Logo" width="600" />
</p>
<div align="center">
<p align="center">Rembg is a tool to remove image backgrounds. It can be used as a CLI, Python library, HTTP server, or Docker container.</p>
<div style="display: flex; flex-direction: row; justify-content: center; gap: 8px; flex-wrap: wrap; margin-top: 8px;">
<a href="https://img.shields.io/badge/License-MIT-blue.svg"><img src="https://img.shields.io/badge/License-MIT-blue.svg" alt="License" /></a>
<a href="https://huggingface.co/spaces/KenjieDec/RemBG"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face-Spaces-blue" alt="Hugging Face Spaces" /></a>
<a href="https://bgremoval.streamlit.app/"><img src="https://img.shields.io/badge/🎈%20Streamlit%20Community-Cloud-blue" alt="Streamlit App" /></a>
<a href="https://colab.research.google.com/github/danielgatis/rembg/blob/main/rembg.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open in Colab" /></a>
<a href="https://repomapr.com/danielgatis/rembg"><img src="https://img.shields.io/badge/RepoMapr-View_Interactive_Diagram-blue?style=flat&logo=github" alt="RepoMapr" /></a>
</div>
</div>
<br/>
<p align="center">
<a href="https://trendshift.io/repositories/2846" target="_blank">
<img src="https://trendshift.io/api/badge/repositories/2846" alt="danielgatis%2Frembg | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
</a>
</p>
## Sponsors
<table>
<tr>
<td align="center" vertical-align="center">
<a href="https://photoroom.com/api/remove-background?utm_source=rembg&utm_medium=github_webpage&utm_campaign=sponsor" >
<img src="https://font-cdn.photoroom.com/media/api-logo.png" width="120px;" alt="Unsplash" />
</a>
</td>
<td align="center" vertical-align="center">
<b>PhotoRoom Remove Background API</b>
<br />
<a href="https://photoroom.com/api/remove-background?utm_source=rembg&utm_medium=github_webpage&utm_campaign=sponsor">https://photoroom.com/api</a>
<br />
<p width="200px">
Fast and accurate background remover API<br/>
</p>
</td>
</tr>
</table>
**If this project has helped you, please consider making a [donation](https://www.buymeacoffee.com/danielgatis).**
## Requirements
```text
python: >=3.11, <3.14
```
## Installation
Choose **one** of the following backends based on your hardware:
### CPU support
```bash
pip install "rembg[cpu]" # for library
pip install "rembg[cpu,cli]" # for library + cli
```
### GPU support (NVIDIA/CUDA)
First, check if your system supports `onnxruntime-gpu` by visiting [onnxruntime.ai](https://onnxruntime.ai/getting-started) and reviewing the installation matrix.
<p style="display: flex;align-items: center;justify-content: center;">
<img alt="onnxruntime-installation-matrix" src="./onnxruntime-installation-matrix.png" width="400" />
</p>
If your system is compatible, run:
```bash
pip install "rembg[gpu]" # for library
pip install "rembg[gpu,cli]" # for library + cli
```
> **Note:** NVIDIA GPUs may require `onnxruntime-gpu`, CUDA, and `cudnn-devel`. See [#668](https://github.com/danielgatis/rembg/issues/668#issuecomment-2689830314) for details. If `rembg[gpu]` doesn't work and you can't install CUDA or `cudnn-devel`, use `rembg[cpu]` with `onnxruntime` instead.
### GPU support (AMD/ROCm)
ROCm support requires the `onnxruntime-rocm` package. Install it by following [AMD's documentation](https://rocm.docs.amd.com/projects/radeon/en/latest/docs/install/native_linux/install-onnx.html).
Once `onnxruntime-rocm` is installed and working, install rembg with ROCm support:
```bash
pip install "rembg[rocm]" # for library
pip install "rembg[rocm,cli]" # for library + cli
```
## Usage as a CLI
After installation, you can use rembg by typing `rembg` in your terminal.
The `rembg` command has 4 subcommands, one for each input type:
- `i` - single files
- `p` - folders (batch processing)
- `s` - HTTP server
- `b` - RGB24 pixel binary stream
You can get help about the main command using:
```shell
rembg --help
```
You can also get help for any subcommand:
```shell
rembg <COMMAND> --help
```
### rembg `i`
Used for processing single files.
**Remove background from a remote image:**
```shell
curl -s http://input.png | rembg i > output.png
```
**Remove background from a local file:**
```shell
rembg i path/to/input.png path/to/output.png
```
**Specify a model:**
```shell
rembg i -m u2netp path/to/input.png path/to/output.png
```
**Return only the mask:**
```shell
rembg i -om path/to/input.png path/to/output.png
```
**Apply alpha matting:**
```shell
rembg i -a path/to/input.png path/to/output.png
```
**Pass extra parameters (SAM example):**
```shell
rembg i -m sam -x '{ "sam_prompt": [{"type": "point", "data": [724, 740], "label": 1}] }' examples/plants-1.jpg examples/plants-1.out.png
```
**Pass extra parameters (custom model):**
```shell
rembg i -m u2net_custom -x '{"model_path": "~/.u2net/u2net.onnx"}' path/to/input.png path/to/output.png
```
### rembg `p`
Used for batch processing entire folders.
**Process all images in a folder:**
```shell
rembg p path/to/input path/to/output
```
**Watch mode (process new/changed files automatically):**
```shell
rembg p -w path/to/input path/to/output
```
### rembg `s`
Used to start an HTTP server.
```shell
rembg s --host 0.0.0.0 --port 7000 --log_level info
```
For complete API documentation, visit: `http://localhost:7000/api`
**Remove background from an image URL:**
```shell
curl -s "http://localhost:7000/api/remove?url=http://input.png" -o output.png
```
**Remove background from an uploaded image:**
```shell
curl -s -F file=@/path/to/input.jpg "http://localhost:7000/api/remove" -o output.png
```
### rembg `b`
Process a sequence of RGB24 images from stdin. This is intended to be used with programs like FFmpeg that output RGB24 pixel data to stdout.
```shell
rembg b <width> <height> -o <output_specifier>
```
**Arguments:**
| Argument | Description |
|----------|-------------|
| `width` | Width of input image(s) |
| `height` | Height of input image(s) |
| `output_specifier` | Printf-style specifier for output filenames (e.g., `output-%03u.png` produces `output-000.png`, `output-001.png`, etc.). Omit to write to stdout. |
**Example with FFmpeg:**
```shell
ffmpeg -i input.mp4 -ss 10 -an -f rawvideo -pix_fmt rgb24 pipe:1 | rembg b 1280 720 -o folder/output-%03u.png
```
> **Note:** The width and height must match FFmpeg's output dimensions. The flags `-an -f rawvideo -pix_fmt rgb24 pipe:1` are required for FFmpeg compatibility.
## Usage as a Library
**Input and output as bytes:**
```python
from rembg import remove
with open('input.png', 'rb') as i:
with open('output.png', 'wb') as o:
input = i.read()
output = remove(input)
o.write(output)
```
**Input and output as a PIL image:**
```python
from rembg import remove
from PIL import Image
input = Image.open('input.png')
output = remove(input)
output.save('output.png')
```
**Input and output as a NumPy array:**
```python
from rembg import remove
import cv2
input = cv2.imread('input.png')
output = remove(input)
cv2.imwrite('output.png', output)
```
**Force output as bytes:**
```python
from rembg import remove
with open('input.png', 'rb') as i:
with open('output.png', 'wb') as o:
input = i.read()
output = remove(input, force_return_bytes=True)
o.write(output)
```
**Batch processing with session reuse (recommended for performance):**
```python
from pathlib import Path
from rembg import remove, new_session
session = new_session()
for file in Path('path/to/folder').glob('*.png'):
input_path = str(file)
output_path = str(file.parent / (file.stem + ".out.png"))
with open(input_path, 'rb') as i:
with open(output_path, 'wb') as o:
input = i.read()
output = remove(input, session=session)
o.write(output)
```
For more examples, see the [examples](USAGE.md) page.
## Usage with Docker
### CPU Only
Replace the `rembg` command with `docker run danielgatis/rembg`:
```shell
docker run -v .:/data danielgatis/rembg i /data/input.png /data/output.png
```
### NVIDIA CUDA GPU Acceleration
**Requirements:** Your host must have the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) installed.
CUDA acceleration requires `cudnn-devel`, so you need to build the Docker image yourself. See [#668](https://github.com/danielgatis/rembg/issues/668#issuecomment-2689914205) for details.
**Build the image:**
```shell
docker build -t rembg-nvidia-cuda-cudnn-gpu -f Dockerfile_nvidia_cuda_cudnn_gpu .
```
> **Note:** This image requires ~11GB of disk space (CPU version is ~1.6GB). Models are not included.
**Run the container:**
```shell
sudo docker run --rm -it --gpus all -v /dev/dri:/dev/dri -v $PWD:/data rembg-nvidia-cuda-cudnn-gpu i -m birefnet-general /data/input.png /data/output.png
```
**Tips:**
- You can create your own NVIDIA CUDA image and install `rembg[gpu,cli]` in it.
- Use `-v /path/to/models/:/root/.u2net` to store model files outside the container, avoiding re-downloads.
## Models
All models are automatically downloaded and saved to `~/.u2net/` on first use.
### Available Models
- u2net ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for general use cases.
- u2netp ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx), [source](https://github.com/xuebinqin/U-2-Net)): A lightweight version of u2net model.
- u2net_human_seg ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for human segmentation.
- u2net_cloth_seg ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx), [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.
- silueta ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx), [source](https://github.com/xuebinqin/U-2-Net/issues/295)): Same as u2net but the size is reduced to 43Mb.
- isnet-general-use ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx), [source](https://github.com/xuebinqin/DIS)): A new pre-trained model for general use cases.
- isnet-anime ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-anime.onnx), [source](https://github.com/SkyTNT/anime-segmentation)): A high-accuracy segmentation for anime character.
- sam ([download encoder](https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-encoder-quant.onnx), [download decoder](https://github.com/danielgatis/rembg/releases/download/v0.0.0/vit_b-decoder-quant.onnx), [source](https://github.com/facebookresearch/segment-anything)): A pre-trained model for any use cases.
- birefnet-general ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx), [source](https://github.com/ZhengPeng7/BiRefNet)): A pre-trained model for general use cases.
- birefnet-general-lite ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx), [source](https://github.com/ZhengPeng7/BiRefNet)): A light pre-trained model for general use cases.
- birefnet-portrait ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx), [source](https://github.com/ZhengPeng7/BiRefNet)): A pre-trained model for human portraits.
- birefnet-dis ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx), [source](https://github.com/ZhengPeng7/BiRefNet)): A pre-trained model for dichotomous image segmentation (DIS).
- birefnet-hrsod ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx), [source](https://github.com/ZhengPeng7/BiRefNet)): A pre-trained model for high-resolution salient object detection (HRSOD).
- birefnet-cod ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx), [source](https://github.com/ZhengPeng7/BiRefNet)): A pre-trained model for concealed object detection (COD).
- birefnet-massive ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx), [source](https://github.com/ZhengPeng7/BiRefNet)): A pre-trained model with massive dataset.
- bria-rmbg ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/bria-rmbg-2.0.onnx), [source](https://huggingface.co/briaai/RMBG-2.0)): A state-of-the-art background removal model by BRIA AI.
## Environment Variables
| Variable | Description |
|----------|-------------|
| `U2NET_HOME` | Path to the directory where models are stored. Defaults to `$XDG_DATA_HOME/.u2net` (or `~/.u2net` if `XDG_DATA_HOME` is not set). |
| `XDG_DATA_HOME` | Base data directory used when `U2NET_HOME` is not set. Defaults to `~`. |
| `MODEL_CHECKSUM_DISABLED` | When set (e.g. `MODEL_CHECKSUM_DISABLED=1`), disables hash verification for downloaded models. This is useful if you want to use your own custom/converted model files without rembg re-downloading the originals. |
| `OMP_NUM_THREADS` | Sets the number of threads used by ONNX Runtime for inference. |
### Using custom model files
If you need to use a modified version of a model (e.g. converted to a different ONNX IR version for compatibility with an older CUDA toolkit), you can prevent rembg from overwriting it:
1. Set `MODEL_CHECKSUM_DISABLED=1`
2. Place your custom `.onnx` file in the models directory (`~/.u2net/` by default) with the expected filename (e.g. `u2net.onnx`)
3. Rembg will detect the file exists and use it without re-downloading
## FAQ
### When will this library support Python version 3.xx?
This library depends on [onnxruntime](https://pypi.org/project/onnxruntime). Python version support is determined by onnxruntime's compatibility.
## Support
If you find this project useful, consider buying me a coffee (or 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>
## Star History
[](https://star-history.com/#danielgatis/rembg&Date)
## License
Copyright (c) 2020-present [Daniel Gatis](https://github.com/danielgatis)
Licensed under the [MIT License](./LICENSE.txt).
================================================
FILE: USAGE.md
================================================
# How to use the remove function
## Load the Image
```python
from PIL import Image
from rembg import new_session, remove
input_path = 'input.png'
output_path = 'output.png'
input = Image.open(input_path)
```
## Removing the background
### Without additional arguments
This defaults to the `u2net` model.
```python
output = remove(input)
output.save(output_path)
```
### With a specific model
You can use the `new_session` function to create a session with a specific model.
```python
model_name = "isnet-general-use"
session = new_session(model_name)
output = remove(input, session=session)
```
### For processing multiple image files
By default, `remove` initialises a new session every call. This can be a large bottleneck if you're having to process multiple images. Initialise a session and pass it in to the `remove` function for fast multi-image support
```python
model_name = "unet"
rembg_session = new_session(model_name)
for img in images:
output = remove(img, session=rembg_session)
```
### With alpha matting
Alpha matting is a post processing step that can be used to improve the quality of the output.
```python
output = remove(input, alpha_matting=True, alpha_matting_foreground_threshold=270,alpha_matting_background_threshold=20, alpha_matting_erode_size=11)
```
### Only mask
If you only want the mask, you can use the `only_mask` argument.
```python
output = remove(input, only_mask=True)
```
### With post processing
You can use the `post_process_mask` argument to post process the mask to get better results.
```python
output = remove(input, post_process_mask=True)
```
### Replacing the background color
You can use the `bgcolor` argument to replace the background color.
```python
output = remove(input, bgcolor=(255, 255, 255, 255))
```
### Using input points
You can use the `input_points` and `input_labels` arguments to specify the points that should be used for the masks. This only works with the `sam` model.
```python
import numpy as np
# Define the points and labels
# The points are defined as [y, x]
input_points = np.array([[400, 350], [700, 400], [200, 400]])
input_labels = np.array([1, 1, 2])
image = remove(image,session=session, input_points=input_points, input_labels=input_labels)
```
## Save the image
```python
output.save(output_path)
```
================================================
FILE: _build-exe.ps1
================================================
# Install Poetry if not already installed
if (-not (Get-Command poetry -ErrorAction SilentlyContinue)) {
pip install poetry
}
# Build CPU version
Write-Host "Building CPU version..." -ForegroundColor Cyan
poetry install --extras "cli cpu"
poetry run pip install pyinstaller
poetry run pyinstaller rembg.spec
Rename-Item -Path "dist/rembg" -NewName "rembg-cpu"
# Build GPU version
Write-Host "Building GPU version..." -ForegroundColor Cyan
poetry install --extras "cli gpu"
poetry run pip install pyinstaller
poetry run pyinstaller rembg.spec --noconfirm
Rename-Item -Path "dist/rembg" -NewName "rembg-gpu"
Write-Host "Build complete!" -ForegroundColor Green
Write-Host "CPU version: dist/rembg-cpu"
Write-Host "GPU version: dist/rembg-gpu"
================================================
FILE: _modpath.iss
================================================
// ----------------------------------------------------------------------------
//
// Inno Setup Ver: 5.4.2
// Script Version: 1.4.2
// Author: Jared Breland <jbreland@legroom.net>
// Homepage: http://www.legroom.net/software
// License: GNU Lesser General Public License (LGPL), version 3
// http://www.gnu.org/licenses/lgpl.html
//
// Script Function:
// Allow modification of environmental path directly from Inno Setup installers
//
// Instructions:
// Copy modpath.iss to the same directory as your setup script
//
// Add this statement to your [Setup] section
// ChangesEnvironment=true
//
// Add this statement to your [Tasks] section
// You can change the Description or Flags
// You can change the Name, but it must match the ModPathName setting below
// Name: modifypath; Description: &Add application directory to your environmental path; Flags: unchecked
//
// Add the following to the end of your [Code] section
// ModPathName defines the name of the task defined above
// ModPathType defines whether the 'user' or 'system' path will be modified;
// this will default to user if anything other than system is set
// setArrayLength must specify the total number of dirs to be added
// Result[0] contains first directory, Result[1] contains second, etc.
// const
// ModPathName = 'modifypath';
// ModPathType = 'user';
//
// function ModPathDir(): TArrayOfString;
// begin
// setArrayLength(Result, 1);
// Result[0] := ExpandConstant('{app}');
// end;
// #include "modpath.iss"
// ----------------------------------------------------------------------------
procedure ModPath();
var
oldpath: String;
newpath: String;
updatepath: Boolean;
pathArr: TArrayOfString;
aExecFile: String;
aExecArr: TArrayOfString;
i, d: Integer;
pathdir: TArrayOfString;
regroot: Integer;
regpath: String;
begin
// Get constants from main script and adjust behavior accordingly
// ModPathType MUST be 'system' or 'user'; force 'user' if invalid
if ModPathType = 'system' then begin
regroot := HKEY_LOCAL_MACHINE;
regpath := 'SYSTEM\CurrentControlSet\Control\Session Manager\Environment';
end else begin
regroot := HKEY_CURRENT_USER;
regpath := 'Environment';
end;
// Get array of new directories and act on each individually
pathdir := ModPathDir();
for d := 0 to GetArrayLength(pathdir)-1 do begin
updatepath := true;
// Modify WinNT path
if UsingWinNT() = true then begin
// Get current path, split into an array
RegQueryStringValue(regroot, regpath, 'Path', oldpath);
oldpath := oldpath + ';';
i := 0;
while (Pos(';', oldpath) > 0) do begin
SetArrayLength(pathArr, i+1);
pathArr[i] := Copy(oldpath, 0, Pos(';', oldpath)-1);
oldpath := Copy(oldpath, Pos(';', oldpath)+1, Length(oldpath));
i := i + 1;
// Check if current directory matches app dir
if pathdir[d] = pathArr[i-1] then begin
// if uninstalling, remove dir from path
if IsUninstaller() = true then begin
continue;
// if installing, flag that dir already exists in path
end else begin
updatepath := false;
end;
end;
// Add current directory to new path
if i = 1 then begin
newpath := pathArr[i-1];
end else begin
newpath := newpath + ';' + pathArr[i-1];
end;
end;
// Append app dir to path if not already included
if (IsUninstaller() = false) AND (updatepath = true) then
newpath := newpath + ';' + pathdir[d];
// Write new path
RegWriteStringValue(regroot, regpath, 'Path', newpath);
// Modify Win9x path
end else begin
// Convert to shortened dirname
pathdir[d] := GetShortName(pathdir[d]);
// If autoexec.bat exists, check if app dir already exists in path
aExecFile := 'C:\AUTOEXEC.BAT';
if FileExists(aExecFile) then begin
LoadStringsFromFile(aExecFile, aExecArr);
for i := 0 to GetArrayLength(aExecArr)-1 do begin
if IsUninstaller() = false then begin
// If app dir already exists while installing, skip add
if (Pos(pathdir[d], aExecArr[i]) > 0) then
updatepath := false;
break;
end else begin
// If app dir exists and = what we originally set, then delete at uninstall
if aExecArr[i] = 'SET PATH=%PATH%;' + pathdir[d] then
aExecArr[i] := '';
end;
end;
end;
// If app dir not found, or autoexec.bat didn't exist, then (create and) append to current path
if (IsUninstaller() = false) AND (updatepath = true) then begin
SaveStringToFile(aExecFile, #13#10 + 'SET PATH=%PATH%;' + pathdir[d], True);
// If uninstalling, write the full autoexec out
end else begin
SaveStringsToFile(aExecFile, aExecArr, False);
end;
end;
end;
end;
// Split a string into an array using passed delimiter
procedure MPExplode(var Dest: TArrayOfString; Text: String; Separator: String);
var
i: Integer;
begin
i := 0;
repeat
SetArrayLength(Dest, i+1);
if Pos(Separator,Text) > 0 then begin
Dest[i] := Copy(Text, 1, Pos(Separator, Text)-1);
Text := Copy(Text, Pos(Separator,Text) + Length(Separator), Length(Text));
i := i + 1;
end else begin
Dest[i] := Text;
Text := '';
end;
until Length(Text)=0;
end;
procedure CurStepChanged(CurStep: TSetupStep);
var
taskname: String;
begin
taskname := ModPathName;
if CurStep = ssPostInstall then
if IsTaskSelected(taskname) then
ModPath();
end;
procedure CurUninstallStepChanged(CurUninstallStep: TUninstallStep);
var
aSelectedTasks: TArrayOfString;
i: Integer;
taskname: String;
regpath: String;
regstring: String;
appid: String;
begin
// only run during actual uninstall
if CurUninstallStep = usUninstall then begin
// get list of selected tasks saved in registry at install time
appid := '{#emit SetupSetting("AppId")}';
if appid = '' then appid := '{#emit SetupSetting("AppName")}';
regpath := ExpandConstant('Software\Microsoft\Windows\CurrentVersion\Uninstall\'+appid+'_is1');
RegQueryStringValue(HKLM, regpath, 'Inno Setup: Selected Tasks', regstring);
if regstring = '' then RegQueryStringValue(HKCU, regpath, 'Inno Setup: Selected Tasks', regstring);
// check each task; if matches modpath taskname, trigger patch removal
if regstring <> '' then begin
taskname := ModPathName;
MPExplode(aSelectedTasks, regstring, ',');
if GetArrayLength(aSelectedTasks) > 0 then begin
for i := 0 to GetArrayLength(aSelectedTasks)-1 do begin
if comparetext(aSelectedTasks[i], taskname) = 0 then
ModPath();
end;
end;
end;
end;
end;
function NeedRestart(): Boolean;
var
taskname: String;
begin
taskname := ModPathName;
if IsTaskSelected(taskname) and not UsingWinNT() then begin
Result := True;
end else begin
Result := False;
end;
end;
================================================
FILE: _setup-cpu.iss
================================================
#define MyAppName "Rembg CPU"
#define MyAppVersion "STABLE"
#define MyAppPublisher "danielgatis"
#define MyAppURL "https://github.com/danielgatis/rembg"
#define MyAppExeName "rembg.exe"
#define MyAppId "49AB7484-212F-4B31-A49F-533A480F3FD4"
[Setup]
AppId={#MyAppId}
AppName={#MyAppName}
AppVersion={#MyAppVersion}
AppPublisher={#MyAppPublisher}
AppPublisherURL={#MyAppURL}
AppSupportURL={#MyAppURL}
AppUpdatesURL={#MyAppURL}
DefaultDirName={autopf}\Rembg
DefaultGroupName=Rembg
DisableProgramGroupPage=yes
OutputBaseFilename=rembg-cli-cpu-installer
Compression=lzma
SolidCompression=yes
WizardStyle=modern
OutputDir=dist
ChangesEnvironment=yes
[Languages]
Name: "english"; MessagesFile: "compiler:Default.isl"
[Files]
Source: "{#SourcePath}dist\rembg-cpu\{#MyAppExeName}"; DestDir: "{app}"; Flags: ignoreversion
Source: "{#SourcePath}dist\rembg-cpu\*"; DestDir: "{app}"; Flags: ignoreversion recursesubdirs createallsubdirs
[Tasks]
Name: modifypath; Description: "Add to PATH variable"
[Icons]
Name: "{group}\Rembg"; Filename: "{app}\{#MyAppExeName}"
[Code]
const
ModPathName = 'modifypath';
ModPathType = 'user';
function ModPathDir(): TArrayOfString;
begin
setArrayLength(Result, 1)
Result[0] := ExpandConstant('{app}');
end;
#include "_modpath.iss"
================================================
FILE: _setup-gpu.iss
================================================
#define MyAppName "Rembg GPU"
#define MyAppVersion "STABLE"
#define MyAppPublisher "danielgatis"
#define MyAppURL "https://github.com/danielgatis/rembg"
#define MyAppExeName "rembg.exe"
#define MyAppId "49AB7484-212F-4B31-A49F-533A480F3FD4"
[Setup]
AppId={#MyAppId}
AppName={#MyAppName}
AppVersion={#MyAppVersion}
AppPublisher={#MyAppPublisher}
AppPublisherURL={#MyAppURL}
AppSupportURL={#MyAppURL}
AppUpdatesURL={#MyAppURL}
DefaultDirName={autopf}\Rembg
DefaultGroupName=Rembg
DisableProgramGroupPage=yes
OutputBaseFilename=rembg-cli-gpu-installer
Compression=lzma
SolidCompression=yes
WizardStyle=modern
OutputDir=dist
ChangesEnvironment=yes
[Languages]
Name: "english"; MessagesFile: "compiler:Default.isl"
[Files]
Source: "{#SourcePath}dist\rembg-gpu\{#MyAppExeName}"; DestDir: "{app}"; Flags: ignoreversion
Source: "{#SourcePath}dist\rembg-gpu\*"; DestDir: "{app}"; Flags: ignoreversion recursesubdirs createallsubdirs
[Tasks]
Name: modifypath; Description: "Add to PATH variable"
[Icons]
Name: "{group}\Rembg"; Filename: "{app}\{#MyAppExeName}"
[Code]
const
ModPathName = 'modifypath';
ModPathType = 'user';
function ModPathDir(): TArrayOfString;
begin
setArrayLength(Result, 1)
Result[0] := ExpandConstant('{app}');
end;
#include "_modpath.iss"
================================================
FILE: docker-compose.yml
================================================
---
# You can set variables in .env file in root folder
#
# PUBLIC_PORT=7000:7000
# REPLICAS_COUNT=1
services:
app:
build: .
command: ["s"]
deploy:
replicas: ${REPLICAS_COUNT:-1}
ports:
- ${PUBLIC_PORT:-7000:7000}
version: '3'
================================================
FILE: man/rembg.1
================================================
.TH REMBG 1 "Januar 2026" "2.0.72" "User Commands"
.SH NAME
rembg \- tool to remove background from images
.SH SYNOPSIS
.B rembg
[OPTIONS] COMMAND [ARGS]...
.SH DESCRIPTION
.B rembg
is a tool to remove images background.
.PP
It works as a command line interface and a library.
.SH OPTIONS
.TP
.BR \-\-version
Show the version and exit.
.TP
.BR \-\-help
Show this message and exit.
.SH SEE ALSO
Full documentation at: <https://github.com/danielgatis/rembg>
================================================
FILE: pyproject.toml
================================================
[tool.poetry]
name = "rembg"
version = "0.0.0" # Managed by poetry-dynamic-versioning
description = "Remove image background"
authors = ["Daniel Gatis <danielgatis@gmail.com>"]
license = "MIT"
readme = "README.md"
homepage = "https://github.com/danielgatis/rembg"
repository = "https://github.com/danielgatis/rembg"
keywords = ["remove", "background", "u2net"]
classifiers = [
"License :: OSI Approved :: MIT License",
"Topic :: Scientific/Engineering",
"Topic :: Scientific/Engineering :: Mathematics",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Software Development",
"Topic :: Software Development :: Libraries",
"Topic :: Software Development :: Libraries :: Python Modules",
"Programming Language :: Python",
"Programming Language :: Python :: 3 :: Only",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
]
packages = [{include = "rembg"}]
[tool.poetry.dependencies]
python = "^3.11"
jsonschema = "^4.25.1"
numpy = "^2.3.0"
pillow = "^12.1.0"
pooch = "^1.8.2"
pymatting = "^1.1.14"
scikit-image = "^0.26.0"
scipy = "^1.16.3"
tqdm = "^4.67.1"
# CPU backend (optional)
onnxruntime = {version = "^1.23.2", optional = true}
# GPU backend (optional) - only available on Linux/Windows, not on macOS
onnxruntime-gpu = {version = "^1.23.2", optional = true, markers = "sys_platform != 'darwin'"}
# ROCm backend (optional) - only available on Linux (latest is 1.22.x)
onnxruntime-rocm = {version = "^1.22.0", optional = true, markers = "sys_platform == 'linux'"}
# CLI dependencies (optional)
aiohttp = {version = "^3.13.2", optional = true}
asyncer = {version = "^0.0.12", optional = true}
click = {version = "^8.3.1", optional = true}
fastapi = {version = "^0.128.0", optional = true}
filetype = {version = "^1.2.0", optional = true}
gradio = {version = "^6.2.0", optional = true}
python-multipart = {version = "^0.0.21", optional = true}
sniffio = {version = "^1.3.1", optional = true}
uvicorn = {version = "^0.40.0", optional = true}
watchdog = {version = "^6.0.0", optional = true}
# Dev dependencies (optional, for pip install .[dev])
bandit = {version = "^1.9.2", optional = true}
black = {version = "^25.12.0", optional = true}
flake8 = {version = "^7.3.0", optional = true}
imagehash = {version = "^4.3.2", optional = true}
isort = {version = "^7.0.0", optional = true}
mypy = {version = "^1.19.1", optional = true}
pytest = {version = "^9.0.2", optional = true}
[tool.poetry.group.dev.dependencies]
bandit = "^1.9.2"
black = "^25.12.0"
flake8 = "^7.3.0"
imagehash = "^4.3.2"
isort = "^7.0.0"
mypy = "^1.19.1"
pytest = "^9.0.2"
[tool.poetry.extras]
cpu = ["onnxruntime"]
gpu = ["onnxruntime-gpu"]
rocm = ["onnxruntime-rocm"]
cli = ["aiohttp", "asyncer", "click", "fastapi", "filetype", "gradio", "python-multipart", "sniffio", "uvicorn", "watchdog"]
dev = ["bandit", "black", "flake8", "imagehash", "isort", "mypy", "pytest"]
[tool.poetry.scripts]
rembg = "rembg.cli:main"
[build-system]
requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
build-backend = "poetry_dynamic_versioning.backend"
[tool.poetry-dynamic-versioning]
enable = true
vcs = "git"
style = "pep440"
pattern = "^v(?P<base>\\d+\\.\\d+\\.\\d+)"
[tool.poetry-dynamic-versioning.substitution]
files = ["rembg/__init__.py"]
================================================
FILE: pytest.ini
================================================
[pytest]
filterwarnings =
ignore::DeprecationWarning
================================================
FILE: rembg/__init__.py
================================================
try:
from importlib.metadata import PackageNotFoundError, version
try:
__version__ = version("rembg")
except PackageNotFoundError:
__version__ = "0.0.0" # Fallback for development
except ImportError:
__version__ = "0.0.0" # Fallback for older Python versions
from .bg import remove
from .session_factory import new_session
================================================
FILE: rembg/bg.py
================================================
import io
import sys
from enum import Enum
from typing import Any, List, Optional, Tuple, Union, cast
import numpy as np
try:
import onnxruntime as ort # type: ignore[import-untyped]
except ImportError:
print("No onnxruntime backend found.")
print("Please install rembg with CPU or GPU support:")
print()
print(' pip install "rembg[cpu]" # for CPU')
print(' pip install "rembg[gpu]" # for NVIDIA/CUDA GPU')
print()
print(
"For more information, see: https://github.com/danielgatis/rembg#installation"
)
sys.exit(1)
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
from pymatting.util.util import stack_images
from scipy.ndimage import binary_erosion, gaussian_filter
from skimage.morphology import disk, opening
from .session_factory import new_session
from .sessions import sessions, sessions_names
from .sessions.base import BaseSession
ort.set_default_logger_severity(3)
kernel = disk(1)
class ReturnType(Enum):
BYTES = 0
PILLOW = 1
NDARRAY = 2
def alpha_matting_cutout(
img: PILImage,
mask: PILImage,
foreground_threshold: int,
background_threshold: int,
erode_structure_size: int,
) -> PILImage:
"""
Perform alpha matting on an image using a given mask and threshold values.
This function takes a PIL image `img` and a PIL image `mask` as input, along with
the `foreground_threshold` and `background_threshold` values used to determine
foreground and background pixels. The `erode_structure_size` parameter specifies
the size of the erosion structure to be applied to the mask.
The function returns a PIL image representing the cutout of the foreground object
from the original image.
"""
if img.mode == "RGBA" or img.mode == "CMYK":
img = img.convert("RGB")
img_array = np.asarray(img)
mask_array = np.asarray(mask)
is_foreground = mask_array > foreground_threshold
is_background = mask_array < background_threshold
structure = None
if erode_structure_size > 0:
structure = np.ones(
(erode_structure_size, erode_structure_size), dtype=np.uint8
)
is_foreground = binary_erosion(is_foreground, structure=structure)
is_background = binary_erosion(is_background, structure=structure, border_value=1)
trimap = np.full(mask_array.shape, dtype=np.uint8, fill_value=128)
trimap[is_foreground] = 255
trimap[is_background] = 0
img_normalized = img_array / 255.0
trimap_normalized = trimap / 255.0
alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
foreground = estimate_foreground_ml(img_normalized, alpha)
cutout = stack_images(foreground, alpha)
cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)
cutout = Image.fromarray(cutout)
return cutout
def naive_cutout(img: PILImage, mask: PILImage) -> PILImage:
"""
Perform a simple cutout operation on an image using a mask.
This function takes a PIL image `img` and a PIL image `mask` as input.
It uses the mask to create a new image where the pixels from `img` are
cut out based on the mask.
The function returns a PIL image representing the cutout of the original
image using the mask.
"""
empty = Image.new("RGBA", (img.size), 0)
cutout = Image.composite(img, empty, mask)
return cutout
def putalpha_cutout(img: PILImage, mask: PILImage) -> PILImage:
"""
Apply the specified mask to the image as an alpha cutout.
Args:
img (PILImage): The image to be modified.
mask (PILImage): The mask to be applied.
Returns:
PILImage: The modified image with the alpha cutout applied.
"""
img.putalpha(mask)
return img
def get_concat_v_multi(imgs: List[PILImage]) -> PILImage:
"""
Concatenate multiple images vertically.
Args:
imgs (List[PILImage]): The list of images to be concatenated.
Returns:
PILImage: The concatenated image.
"""
pivot = imgs.pop(0)
for im in imgs:
pivot = get_concat_v(pivot, im)
return pivot
def get_concat_v(img1: PILImage, img2: PILImage) -> PILImage:
"""
Concatenate two images vertically.
Args:
img1 (PILImage): The first image.
img2 (PILImage): The second image to be concatenated below the first image.
Returns:
PILImage: The concatenated image.
"""
dst = Image.new("RGBA", (img1.width, img1.height + img2.height))
dst.paste(img1, (0, 0))
dst.paste(img2, (0, img1.height))
return dst
def post_process(mask: np.ndarray) -> np.ndarray:
"""
Post Process the mask for a smooth boundary by applying Morphological Operations
Research based on paper: https://www.sciencedirect.com/science/article/pii/S2352914821000757
args:
mask: Binary Numpy Mask
"""
mask = opening(mask, kernel)
mask = gaussian_filter(mask.astype(np.float64), sigma=2)
mask = np.where(mask < 127, 0, 255).astype(np.uint8)
return mask
def apply_background_color(img: PILImage, color: Tuple[int, int, int, int]) -> PILImage:
"""
Apply the specified background color to the image.
Args:
img (PILImage): The image to be modified.
color (Tuple[int, int, int, int]): The RGBA color to be applied.
Returns:
PILImage: The modified image with the background color applied.
"""
background = Image.new("RGBA", img.size, tuple(color))
colored_image = Image.alpha_composite(background, img)
return colored_image
def fix_image_orientation(img: PILImage) -> PILImage:
"""
Fix the orientation of the image based on its EXIF data.
Args:
img (PILImage): The image to be fixed.
Returns:
PILImage: The fixed image.
"""
return cast(PILImage, ImageOps.exif_transpose(img))
def download_models(models: tuple[str, ...]) -> None:
"""
Download models for image processing.
"""
if len(models) == 0:
print("No models specified, downloading all models")
models = tuple(sessions_names)
for model in models:
session = sessions.get(model)
if session is None:
print(f"Error: no model found: {model}")
sys.exit(1)
else:
print(f"Downloading model: {model}")
try:
session.download_models()
except Exception as e:
print(f"Error downloading model: {e}")
def remove(
data: Union[bytes, PILImage, np.ndarray],
alpha_matting: bool = False,
alpha_matting_foreground_threshold: int = 240,
alpha_matting_background_threshold: int = 10,
alpha_matting_erode_size: int = 10,
session: Optional[BaseSession] = None,
only_mask: bool = False,
post_process_mask: bool = False,
bgcolor: Optional[Tuple[int, int, int, int]] = None,
force_return_bytes: bool = False,
*args: Optional[Any],
**kwargs: Optional[Any],
) -> Union[bytes, PILImage, np.ndarray]:
"""
Remove the background from an input image.
This function takes in various parameters and returns a modified version of the input image with the background removed. The function can handle input data in the form of bytes, a PIL image, or a numpy array. The function first checks the type of the input data and converts it to a PIL image if necessary. It then fixes the orientation of the image and proceeds to perform background removal using the 'u2net' model. The result is a list of binary masks representing the foreground objects in the image. These masks are post-processed and combined to create a final cutout image. If a background color is provided, it is applied to the cutout image. The function returns the resulting cutout image in the format specified by the input 'return_type' parameter or as python bytes if force_return_bytes is true.
Parameters:
data (Union[bytes, PILImage, np.ndarray]): The input image data.
alpha_matting (bool, optional): Flag indicating whether to use alpha matting. Defaults to False.
alpha_matting_foreground_threshold (int, optional): Foreground threshold for alpha matting. Defaults to 240.
alpha_matting_background_threshold (int, optional): Background threshold for alpha matting. Defaults to 10.
alpha_matting_erode_size (int, optional): Erosion size for alpha matting. Defaults to 10.
session (Optional[BaseSession], optional): A session object for the 'u2net' model. Defaults to None.
only_mask (bool, optional): Flag indicating whether to return only the binary masks. Defaults to False.
post_process_mask (bool, optional): Flag indicating whether to post-process the masks. Defaults to False.
bgcolor (Optional[Tuple[int, int, int, int]], optional): Background color for the cutout image. Defaults to None.
force_return_bytes (bool, optional): Flag indicating whether to return the cutout image as bytes. Defaults to False.
*args (Optional[Any]): Additional positional arguments.
**kwargs (Optional[Any]): Additional keyword arguments.
Returns:
Union[bytes, PILImage, np.ndarray]: The cutout image with the background removed.
"""
if isinstance(data, bytes) or force_return_bytes:
return_type = ReturnType.BYTES
img = cast(PILImage, Image.open(io.BytesIO(cast(bytes, data))))
elif isinstance(data, PILImage):
return_type = ReturnType.PILLOW
img = cast(PILImage, data)
elif isinstance(data, np.ndarray):
return_type = ReturnType.NDARRAY
img = cast(PILImage, Image.fromarray(data))
else:
raise ValueError(
"Input type {} is not supported. Try using force_return_bytes=True to force python bytes output".format(
type(data)
)
)
putalpha = kwargs.pop("putalpha", False)
# Fix image orientation
img = fix_image_orientation(img)
if session is None:
session = new_session("u2net", *args, **kwargs)
masks = session.predict(img, *args, **kwargs)
cutouts = []
for mask in masks:
if post_process_mask:
mask = Image.fromarray(post_process(np.array(mask)))
if only_mask:
cutout = mask
elif alpha_matting:
try:
cutout = alpha_matting_cutout(
img,
mask,
alpha_matting_foreground_threshold,
alpha_matting_background_threshold,
alpha_matting_erode_size,
)
except ValueError:
if putalpha:
cutout = putalpha_cutout(img, mask)
else:
cutout = naive_cutout(img, mask)
else:
if putalpha:
cutout = putalpha_cutout(img, mask)
else:
cutout = naive_cutout(img, mask)
cutouts.append(cutout)
cutout = img
if len(cutouts) > 0:
cutout = get_concat_v_multi(cutouts)
if bgcolor is not None and not only_mask:
cutout = apply_background_color(cutout, bgcolor)
if ReturnType.PILLOW == return_type:
return cutout
if ReturnType.NDARRAY == return_type:
return np.asarray(cutout)
bio = io.BytesIO()
cutout.save(bio, "PNG")
bio.seek(0)
return bio.read()
================================================
FILE: rembg/cli.py
================================================
import sys
# Fast path for --version (avoid importing heavy dependencies)
if len(sys.argv) == 2 and sys.argv[1] in ("--version", "-V"):
from importlib.metadata import version
print(f"rembg, version {version('rembg')}")
sys.exit(0)
try:
import click
except ImportError:
print("The CLI dependencies are not installed.")
print("Please install rembg with CLI support:")
print()
print(' pip install "rembg[cpu,cli]" # for CPU')
print(' pip install "rembg[gpu,cli]" # for NVIDIA/CUDA GPU')
print()
print(
"For more information, see: https://github.com/danielgatis/rembg#installation"
)
sys.exit(1)
from . import __version__
from .commands import command_functions
@click.group()
@click.version_option(version=__version__)
def main() -> None:
pass
for command in command_functions:
main.add_command(command)
================================================
FILE: rembg/commands/__init__.py
================================================
command_functions = []
from .b_command import b_command
from .d_command import d_command
from .i_command import i_command
from .p_command import p_command
from .s_command import s_command
command_functions.append(b_command)
command_functions.append(d_command)
command_functions.append(i_command)
command_functions.append(p_command)
command_functions.append(s_command)
================================================
FILE: rembg/commands/b_command.py
================================================
import asyncio
import io
import json
import os
import sys
from typing import IO
import click
import PIL
from ..bg import remove
from ..session_factory import new_session
from ..sessions import sessions_names
@click.command( # type: ignore
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=(0, 0, 0, 0),
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 b_command(
model: str,
extras: str,
image_width: int,
image_height: int,
output_specifier: str,
**kwargs
) -> None:
"""
Command-line interface for processing images by removing the background using a specified model and generating a mask.
This CLI command takes several options and arguments to configure the background removal process and save the processed images.
Parameters:
model (str): The name of the model to use for background removal.
extras (str): Additional options in JSON format that can be passed to customize the background removal process.
image_width (int): The width of the input images in pixels.
image_height (int): The height of the input images in pixels.
output_specifier (str): A printf-style specifier for the output filenames. If specified, the processed images will be saved to the specified output directory with filenames generated using the specifier.
**kwargs: Additional keyword arguments that can be used to customize the background removal process.
Returns:
None
"""
if extras:
try:
kwargs.update(json.loads(extras))
except Exception:
raise click.BadParameter("extras must be a valid JSON string")
session = new_session(model, **kwargs)
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: PIL.Image.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 = PIL.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())
================================================
FILE: rembg/commands/d_command.py
================================================
import click
from ..bg import download_models
@click.command( # type: ignore
name="d",
help="download models",
)
@click.argument("models", nargs=-1)
def d_command(models: tuple[str, ...]) -> None:
"""
Download models
"""
download_models(models)
================================================
FILE: rembg/commands/i_command.py
================================================
import json
import sys
from typing import IO
import click
from ..bg import remove
from ..session_factory import new_session
from ..sessions import sessions_names
@click.command( # type: ignore
name="i",
help="for a file 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=(0, 0, 0, 0),
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.argument(
"input", default=(None if sys.stdin.isatty() else "-"), type=click.File("rb")
)
@click.argument(
"output",
default=(None if sys.stdin.isatty() else "-"),
type=click.File("wb", lazy=True),
)
def i_command(model: str, extras: str, input: IO, output: IO, **kwargs) -> None:
"""
Click command line interface function to process an input file based on the provided options.
This function is the entry point for the CLI program. It reads an input file, applies image processing operations based on the provided options, and writes the output to a file.
Parameters:
model (str): The name of the model to use for image processing.
extras (str): Additional options in JSON format.
input: The input file to process.
output: The output file to write the processed image to.
**kwargs: Additional keyword arguments corresponding to the command line options.
Returns:
None
"""
try:
kwargs.update(json.loads(extras))
except Exception:
pass
output.write(remove(input.read(), session=new_session(model, **kwargs), **kwargs))
================================================
FILE: rembg/commands/p_command.py
================================================
import json
import pathlib
import time
from typing import cast
import click
import filetype
from tqdm import tqdm
from watchdog.events import FileSystemEvent, FileSystemEventHandler
from watchdog.observers import Observer
from ..bg import remove
from ..session_factory import new_session
from ..sessions import sessions_names
@click.command( # type: ignore
name="p",
help="for a folder 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(
"-w",
"--watch",
default=False,
is_flag=True,
show_default=True,
help="watches a folder for changes",
)
@click.option(
"-d",
"--delete_input",
default=False,
is_flag=True,
show_default=True,
help="delete input file after processing",
)
@click.option(
"-bgc",
"--bgcolor",
default=(0, 0, 0, 0),
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.argument(
"input",
type=click.Path(
exists=True,
path_type=pathlib.Path,
file_okay=False,
dir_okay=True,
readable=True,
),
)
@click.argument(
"output",
type=click.Path(
exists=False,
path_type=pathlib.Path,
file_okay=False,
dir_okay=True,
writable=True,
),
)
def p_command(
model: str,
extras: str,
input: pathlib.Path,
output: pathlib.Path,
watch: bool,
delete_input: bool,
**kwargs,
) -> None:
"""
Command-line interface (CLI) program for performing background removal on images in a folder.
This program takes a folder as input and uses a specified model to remove the background from the images in the folder.
It provides various options for configuration, such as choosing the model, enabling alpha matting, setting trimap thresholds, erode size, etc.
Additional options include outputting only the mask and post-processing the mask.
The program can also watch the input folder for changes and automatically process new images.
The resulting images with the background removed are saved in the specified output folder.
Parameters:
model (str): The name of the model to use for background removal.
extras (str): Additional options in JSON format.
input (pathlib.Path): The path to the input folder.
output (pathlib.Path): The path to the output folder.
watch (bool): Whether to watch the input folder for changes.
delete_input (bool): Whether to delete the input file after processing.
**kwargs: Additional keyword arguments.
Returns:
None
"""
try:
kwargs.update(json.loads(extras))
except Exception:
pass
session = new_session(model, **kwargs)
def process(each_input: pathlib.Path) -> None:
try:
mimetype = filetype.guess(each_input)
if mimetype is None:
return
if mimetype.mime.find("image") < 0:
return
each_output = (output / each_input.name).with_suffix(".png")
each_output.parents[0].mkdir(parents=True, exist_ok=True)
if not each_output.exists():
each_output.write_bytes(
cast(
bytes,
remove(each_input.read_bytes(), session=session, **kwargs),
)
)
if watch:
print(
f"processed: {each_input.absolute()} -> {each_output.absolute()}"
)
if delete_input:
each_input.unlink()
except Exception as e:
print(e)
inputs = list(input.glob("**/*"))
inputs_tqdm = inputs if watch else tqdm(inputs)
for each_input in inputs_tqdm:
if not each_input.is_dir():
process(each_input)
if watch:
should_watch = True
observer = Observer()
class EventHandler(FileSystemEventHandler):
def on_any_event(self, event: FileSystemEvent) -> None:
src_path = cast(str, event.src_path)
if (
not (
event.is_directory or event.event_type in ["deleted", "closed"]
)
and pathlib.Path(src_path).exists()
):
if src_path.endswith("stop.txt"):
nonlocal should_watch
should_watch = False
pathlib.Path(src_path).unlink()
return
process(pathlib.Path(src_path))
event_handler = EventHandler()
observer.schedule(event_handler, str(input), recursive=False)
observer.start()
try:
while should_watch:
time.sleep(1)
finally:
observer.stop()
observer.join()
================================================
FILE: rembg/commands/s_command.py
================================================
import json
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
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import Response
from .. import __version__
from ..bg import remove
from ..session_factory import new_session
from ..sessions import sessions_names
from ..sessions.base import BaseSession
@click.command( # type: ignore
name="s",
help="for a http server",
)
@click.option(
"-p",
"--port",
default=7000,
type=int,
show_default=True,
help="port",
)
@click.option(
"-h",
"--host",
default="0.0.0.0",
type=str,
show_default=True,
help="host",
)
@click.option(
"-l",
"--log_level",
default="info",
type=str,
show_default=True,
help="log level",
)
@click.option(
"-t",
"--threads",
default=None,
type=int,
show_default=True,
help="number of worker threads",
)
def s_command(port: int, host: str, log_level: str, threads: int) -> None:
"""
Command-line interface for running the FastAPI web server.
This function starts the FastAPI web server with the specified port and log level.
If the number of worker threads is specified, it sets the thread limiter accordingly.
"""
sessions: dict[str, BaseSession] = {}
tags_metadata = [
{
"name": "Background Removal",
"description": "Endpoints that perform background removal with different image sources.",
"externalDocs": {
"description": "GitHub Source",
"url": "https://github.com/danielgatis/rembg",
},
},
]
app = FastAPI(
title="Rembg",
description="Rembg is a tool to remove images background. That is it.",
version=__version__,
contact={
"name": "Daniel Gatis",
"url": "https://github.com/danielgatis",
"email": "danielgatis@gmail.com",
},
license_info={
"name": "MIT License",
"url": "https://github.com/danielgatis/rembg/blob/main/LICENSE.txt",
},
openapi_tags=tags_metadata,
docs_url="/api",
)
app.add_middleware(
CORSMiddleware,
allow_credentials=True,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
class CommonQueryParams:
def __init__(
self,
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(
default=240,
ge=0,
le=255,
description="Alpha Matting (Foreground Threshold)",
),
ab: int = Query(
default=10,
ge=0,
le=255,
description="Alpha Matting (Background Threshold)",
),
ae: int = Query(
default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
),
om: bool = Query(default=False, description="Only Mask"),
ppm: bool = Query(default=False, description="Post Process Mask"),
bgc: Optional[str] = Query(default=None, description="Background Color"),
extras: Optional[str] = Query(
default=None, description="Extra parameters as JSON"
),
):
self.model = model
self.a = a
self.af = af
self.ab = ab
self.ae = ae
self.om = om
self.ppm = ppm
self.extras = extras
self.bgc = (
cast(Tuple[int, int, int, int], tuple(map(int, bgc.split(","))))
if bgc
else None
)
class CommonQueryPostParams:
def __init__(
self,
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(
default=240,
ge=0,
le=255,
description="Alpha Matting (Foreground Threshold)",
),
ab: int = Form(
default=10,
ge=0,
le=255,
description="Alpha Matting (Background Threshold)",
),
ae: int = Form(
default=10, ge=0, description="Alpha Matting (Erode Structure Size)"
),
om: bool = Form(default=False, description="Only Mask"),
ppm: bool = Form(default=False, description="Post Process Mask"),
bgc: Optional[str] = Query(default=None, description="Background Color"),
extras: Optional[str] = Query(
default=None, description="Extra parameters as JSON"
),
):
self.model = model
self.a = a
self.af = af
self.ab = ab
self.ae = ae
self.om = om
self.ppm = ppm
self.extras = extras
self.bgc = (
cast(Tuple[int, int, int, int], tuple(map(int, bgc.split(","))))
if bgc
else None
)
def im_without_bg(content: bytes, commons: CommonQueryParams) -> Response:
kwargs = {}
if commons.extras:
try:
kwargs.update(json.loads(commons.extras))
except Exception:
pass
session = sessions.get(commons.model)
if session is None:
session = new_session(commons.model, **kwargs)
sessions[commons.model] = session
return Response(
remove(
content,
session=session,
alpha_matting=commons.a,
alpha_matting_foreground_threshold=commons.af,
alpha_matting_background_threshold=commons.ab,
alpha_matting_erode_size=commons.ae,
only_mask=commons.om,
post_process_mask=commons.ppm,
bgcolor=commons.bgc,
**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
RunVar("_default_thread_limiter").set(CapacityLimiter(threads))
@app.get(
path="/api/remove",
tags=["Background Removal"],
summary="Remove from URL",
description="Removes the background from an image obtained by retrieving an URL.",
)
async def get_index(
url: str = Query(
default=..., description="URL of the image that has to be processed."
),
commons: CommonQueryParams = Depends(),
):
async with aiohttp.ClientSession() as session:
async with session.get(url) as response:
file = await response.read()
return await asyncify(im_without_bg)(file, commons)
@app.post(
path="/api/remove",
tags=["Background Removal"],
summary="Remove from Stream",
description="Removes the background from an image sent within the request itself.",
)
async def post_index(
file: bytes = File(
default=...,
description="Image file (byte stream) that has to be processed.",
),
commons: CommonQueryPostParams = Depends(),
):
return await asyncify(im_without_bg)(file, commons) # type: ignore
def gr_app(app):
def inference(input_path, model, *args):
output_path = "output.png"
a, af, ab, ae, om, ppm, cmd_args = args
kwargs = {
"alpha_matting": a,
"alpha_matting_foreground_threshold": af,
"alpha_matting_background_threshold": ab,
"alpha_matting_erode_size": ae,
"only_mask": om,
"post_process_mask": ppm,
}
if cmd_args:
kwargs.update(json.loads(cmd_args))
session = sessions.get(model)
if session is None:
session = new_session(model, **kwargs)
sessions[model] = session
kwargs["session"] = session
with open(input_path, "rb") as i:
with open(output_path, "wb") as o:
input = i.read()
output = remove(input, **kwargs)
o.write(output)
return os.path.join(output_path)
interface = gr.Interface(
inference,
[
gr.components.Image(type="filepath", label="Input"),
gr.components.Dropdown(sessions_names, value="u2net", label="Models"),
gr.components.Checkbox(value=True, label="Alpha matting"),
gr.components.Slider(
value=240, minimum=0, maximum=255, label="Foreground threshold"
),
gr.components.Slider(
value=10, minimum=0, maximum=255, label="Background threshold"
),
gr.components.Slider(
value=40, minimum=0, maximum=255, label="Erosion size"
),
gr.components.Checkbox(value=False, label="Only mask"),
gr.components.Checkbox(value=True, label="Post process mask"),
gr.components.Textbox(label="Arguments"),
],
gr.components.Image(type="filepath", label="Output"),
concurrency_limit=3,
analytics_enabled=False,
)
app = gr.mount_gradio_app(app, interface, path="/")
return app
print(
f"To access the API documentation, go to http://{'localhost' if host == '0.0.0.0' else host}:{port}/api"
)
print(
f"To access the UI, go to http://{'localhost' if host == '0.0.0.0' else host}:{port}"
)
uvicorn.run(gr_app(app), host=host, port=port, log_level=log_level)
================================================
FILE: rembg/session_factory.py
================================================
import os
from typing import Optional, Type
import onnxruntime as ort
from .sessions import sessions_class
from .sessions.base import BaseSession
from .sessions.u2net import U2netSession
def new_session(model_name: str = "u2net", *args, **kwargs) -> BaseSession:
"""
Create a new session object based on the specified model name.
This function searches for the session class based on the model name in the 'sessions_class' list.
It then creates an instance of the session class with the provided arguments.
The 'sess_opts' object is created using the 'ort.SessionOptions()' constructor.
If the 'OMP_NUM_THREADS' environment variable is set, the 'inter_op_num_threads' option of 'sess_opts' is set to its value.
Parameters:
model_name (str): The name of the model.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Raises:
ValueError: If no session class with the given `model_name` is found.
Returns:
BaseSession: The created session object.
"""
session_class: Optional[Type[BaseSession]] = None
for sc in sessions_class:
if sc.name() == model_name:
session_class = sc
break
if session_class is None:
raise ValueError(f"No session class found for model '{model_name}'")
sess_opts = ort.SessionOptions()
if "OMP_NUM_THREADS" in os.environ:
threads = int(os.environ["OMP_NUM_THREADS"])
sess_opts.inter_op_num_threads = threads
sess_opts.intra_op_num_threads = threads
return session_class(model_name, sess_opts, *args, **kwargs)
================================================
FILE: rembg/sessions/__init__.py
================================================
from __future__ import annotations
from typing import Dict, List
from .base import BaseSession
sessions: Dict[str, type[BaseSession]] = {}
from .birefnet_general import BiRefNetSessionGeneral
sessions[BiRefNetSessionGeneral.name()] = BiRefNetSessionGeneral
from .birefnet_general_lite import BiRefNetSessionGeneralLite
sessions[BiRefNetSessionGeneralLite.name()] = BiRefNetSessionGeneralLite
from .birefnet_portrait import BiRefNetSessionPortrait
sessions[BiRefNetSessionPortrait.name()] = BiRefNetSessionPortrait
from .birefnet_dis import BiRefNetSessionDIS
sessions[BiRefNetSessionDIS.name()] = BiRefNetSessionDIS
from .birefnet_hrsod import BiRefNetSessionHRSOD
sessions[BiRefNetSessionHRSOD.name()] = BiRefNetSessionHRSOD
from .birefnet_cod import BiRefNetSessionCOD
sessions[BiRefNetSessionCOD.name()] = BiRefNetSessionCOD
from .birefnet_massive import BiRefNetSessionMassive
sessions[BiRefNetSessionMassive.name()] = BiRefNetSessionMassive
from .dis_anime import DisSession
sessions[DisSession.name()] = DisSession
from .dis_custom import DisCustomSession
sessions[DisCustomSession.name()] = DisCustomSession
from .dis_general_use import DisSession as DisSessionGeneralUse
sessions[DisSessionGeneralUse.name()] = DisSessionGeneralUse
from .sam import SamSession
sessions[SamSession.name()] = SamSession
from .silueta import SiluetaSession
sessions[SiluetaSession.name()] = SiluetaSession
from .u2net_cloth_seg import Unet2ClothSession
sessions[Unet2ClothSession.name()] = Unet2ClothSession
from .u2net_custom import U2netCustomSession
sessions[U2netCustomSession.name()] = U2netCustomSession
from .u2net_human_seg import U2netHumanSegSession
sessions[U2netHumanSegSession.name()] = U2netHumanSegSession
from .u2net import U2netSession
sessions[U2netSession.name()] = U2netSession
from .u2netp import U2netpSession
sessions[U2netpSession.name()] = U2netpSession
from .bria_rmbg import BriaRmBgSession
sessions[BriaRmBgSession.name()] = BriaRmBgSession
from .ben_custom import BenCustomSession
sessions[BenCustomSession.name()] = BenCustomSession
sessions_names = list(sessions.keys())
sessions_class = list(sessions.values())
================================================
FILE: rembg/sessions/base.py
================================================
import os
from typing import Dict, List, Tuple
import numpy as np
import onnxruntime as ort
from PIL import Image
from PIL.Image import Image as PILImage
class BaseSession:
"""This is a base class for managing a session with a machine learning model."""
def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
"""Initialize an instance of the BaseSession class."""
self.model_name = model_name
if "providers" in kwargs and isinstance(kwargs["providers"], list):
providers = kwargs.pop("providers")
else:
device_type = ort.get_device()
if (
device_type == "GPU"
and "CUDAExecutionProvider" in ort.get_available_providers()
):
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
elif (
device_type[0:3] == "GPU"
and "ROCMExecutionProvider" in ort.get_available_providers()
):
providers = ["ROCMExecutionProvider", "CPUExecutionProvider"]
else:
providers = ["CPUExecutionProvider"]
self.inner_session = ort.InferenceSession(
str(self.__class__.download_models(*args, **kwargs)),
sess_options=sess_opts,
providers=providers,
)
def normalize(
self,
img: PILImage,
mean: Tuple[float, float, float],
std: Tuple[float, float, float],
size: Tuple[int, int],
*args,
**kwargs
) -> Dict[str, np.ndarray]:
im = img.convert("RGB").resize(size, Image.Resampling.LANCZOS)
im_ary = np.array(im)
im_ary = im_ary / max(np.max(im_ary), 1e-6)
tmpImg = np.zeros((im_ary.shape[0], im_ary.shape[1], 3))
tmpImg[:, :, 0] = (im_ary[:, :, 0] - mean[0]) / std[0]
tmpImg[:, :, 1] = (im_ary[:, :, 1] - mean[1]) / std[1]
tmpImg[:, :, 2] = (im_ary[:, :, 2] - mean[2]) / std[2]
tmpImg = tmpImg.transpose((2, 0, 1))
return {
self.inner_session.get_inputs()[0]
.name: np.expand_dims(tmpImg, 0)
.astype(np.float32)
}
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(
os.getenv(
"U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net")
)
)
@classmethod
def download_models(cls, *args, **kwargs):
raise NotImplementedError
@classmethod
def name(cls, *args, **kwargs):
raise NotImplementedError
================================================
FILE: rembg/sessions/ben_custom.py
================================================
import os
from typing import List
import numpy as np
import onnxruntime as ort
from PIL import Image
from PIL.Image import Image as PILImage
from .base import BaseSession
class BenCustomSession(BaseSession):
"""This is a class representing a custom session for the Ben model."""
def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
"""
Initialize a new BenCustomSession object.
Parameters:
model_name (str): The name of the model.
sess_opts: The session options.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path is required")
super().__init__(model_name, sess_opts, *args, **kwargs)
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predicts the mask image for the input image.
This method takes a PILImage object as input and returns a list of PILImage objects as output. It performs several image processing operations to generate the mask image.
Parameters:
img (PILImage): The input image.
Returns:
List[PILImage]: A list of PILImage objects representing the generated mask image.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(img, (0.5, 0.5, 0.5), (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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Download the model files.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The absolute path to the model files.
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path is required")
return os.path.abspath(os.path.expanduser(model_path))
@classmethod
def name(cls, *args, **kwargs):
"""
Get the name of the model.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the model.
"""
return "ben_custom"
================================================
FILE: rembg/sessions/birefnet_cod.py
================================================
import os
import pooch
from . import BiRefNetSessionGeneral
class BiRefNetSessionCOD(BiRefNetSessionGeneral):
"""
This class represents a BiRefNet-COD session, which is a subclass of BiRefNetSessionGeneral.
"""
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BiRefNet-COD model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-COD-epoch_125.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:f6d0d21ca89d287f17e7afe9f5fd3b45"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the BiRefNet-COD session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "birefnet-cod"
================================================
FILE: rembg/sessions/birefnet_dis.py
================================================
import os
import pooch
from . import BiRefNetSessionGeneral
class BiRefNetSessionDIS(BiRefNetSessionGeneral):
"""
This class represents a BiRefNet-DIS session, which is a subclass of BiRefNetSessionGeneral.
"""
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BiRefNet-DIS model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-DIS-epoch_590.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:2d4d44102b446f33a4ebb2e56c051f2b"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the BiRefNet-DIS session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "birefnet-dis"
================================================
FILE: rembg/sessions/birefnet_general.py
================================================
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 BiRefNetSessionGeneral(BaseSession):
"""
This class represents a BiRefNet-General session, which is a subclass of BaseSession.
"""
def sigmoid(self, mat):
return 1 / (1 + np.exp(-mat))
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predicts the output masks for the input image using the inner session.
Parameters:
img (PILImage): The input image.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: The list of output masks.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (1024, 1024)
),
)
pred = self.sigmoid(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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BiRefNet-General model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-epoch_244.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:7a35a0141cbbc80de11d9c9a28f52697"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the BiRefNet-General session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "birefnet-general"
================================================
FILE: rembg/sessions/birefnet_general_lite.py
================================================
import os
import pooch
from . import BiRefNetSessionGeneral
class BiRefNetSessionGeneralLite(BiRefNetSessionGeneral):
"""
This class represents a BiRefNet-General-Lite session, which is a subclass of BiRefNetSessionGeneral.
"""
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BiRefNet-General-Lite model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-general-bb_swin_v1_tiny-epoch_232.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:4fab47adc4ff364be1713e97b7e66334"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the BiRefNet-General-Lite session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "birefnet-general-lite"
================================================
FILE: rembg/sessions/birefnet_hrsod.py
================================================
import os
import pooch
from . import BiRefNetSessionGeneral
class BiRefNetSessionHRSOD(BiRefNetSessionGeneral):
"""
This class represents a BiRefNet-HRSOD session, which is a subclass of BiRefNetSessionGeneral.
"""
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BiRefNet-HRSOD model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-HRSOD_DHU-epoch_115.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:c017ade5de8a50ff0fd74d790d268dda"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the BiRefNet-HRSOD session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "birefnet-hrsod"
================================================
FILE: rembg/sessions/birefnet_massive.py
================================================
import os
import pooch
from . import BiRefNetSessionGeneral
class BiRefNetSessionMassive(BiRefNetSessionGeneral):
"""
This class represents a BiRefNet-Massive session, which is a subclass of BiRefNetSessionGeneral.
"""
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BiRefNet-Massive model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-massive-TR_DIS5K_TR_TEs-epoch_420.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:33e726a2136a3d59eb0fdf613e31e3e9"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the BiRefNet-Massive session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "birefnet-massive"
================================================
FILE: rembg/sessions/birefnet_portrait.py
================================================
import os
import pooch
from . import BiRefNetSessionGeneral
class BiRefNetSessionPortrait(BiRefNetSessionGeneral):
"""
This class represents a BiRefNet-Portrait session, which is a subclass of BiRefNetSessionGeneral.
"""
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BiRefNet-Portrait model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/BiRefNet-portrait-epoch_150.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:c3a64a6abf20250d090cd055f12a3b67"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the BiRefNet-Portrait session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "birefnet-portrait"
================================================
FILE: rembg/sessions/bria_rmbg.py
================================================
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 BriaRmBgSession(BaseSession):
"""
This class represents a Bria-rmbg-2.0 session, which is a subclass of BaseSession.
"""
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predicts the output masks for the input image using the inner session.
Parameters:
img (PILImage): The input image.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: The list of output masks.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the BRIA-RMBG 2.0 model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/bria-rmbg-2.0.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "sha256:5b486f08200f513f460da46dd701db5fbb47d79b4be4b708a19444bcd4e79958"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the Bria-rmbg session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "bria-rmbg"
================================================
FILE: rembg/sessions/dis_anime.py
================================================
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):
"""
This class represents a session for object detection.
"""
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Use a pre-trained model to predict the object in the given image.
Parameters:
img (PILImage): The input image.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
List[PILImage]: A list of predicted mask images.
"""
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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Download the pre-trained models.
Parameters:
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The path of the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.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(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Get the name of the pre-trained model.
Parameters:
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The name of the pre-trained model.
"""
return "isnet-anime"
================================================
FILE: rembg/sessions/dis_custom.py
================================================
import os
from typing import List
import numpy as np
import onnxruntime as ort
from PIL import Image
from PIL.Image import Image as PILImage
from .base import BaseSession
class DisCustomSession(BaseSession):
"""This is a class representing a custom session for the Dis model."""
def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
"""
Initialize a new DisCustomSession object.
Parameters:
model_name (str): The name of the model.
sess_opts: The session options.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path is required")
super().__init__(model_name, sess_opts, *args, **kwargs)
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predicts the mask image for the input image.
This method takes a PILImage object as input and returns a list of PILImage objects as output. It performs several image processing operations to generate the mask image.
Parameters:
img (PILImage): The input image.
Returns:
List[PILImage]: A list of PILImage objects representing the generated mask image.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(img, (0.5, 0.5, 0.5), (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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Download the model files.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The absolute path to the model files.
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path is required")
return os.path.abspath(os.path.expanduser(model_path))
@classmethod
def name(cls, *args, **kwargs):
"""
Get the name of the model.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the model.
"""
return "dis_custom"
================================================
FILE: rembg/sessions/dis_general_use.py
================================================
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]:
"""
Predicts the mask image for the input image.
This method takes a PILImage object as input and returns a list of PILImage objects as output. It performs several image processing operations to generate the mask image.
Parameters:
img (PILImage): The input image.
Returns:
List[PILImage]: A list of PILImage objects representing the generated mask image.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(img, (0.5, 0.5, 0.5), (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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the pre-trained model file.
This class method downloads the pre-trained model file from a specified URL using the pooch library.
Parameters:
args: Additional positional arguments.
kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.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(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the model.
This class method returns the name of the model.
Parameters:
args: Additional positional arguments.
kwargs: Additional keyword arguments.
Returns:
str: The name of the model.
"""
return "isnet-general-use"
================================================
FILE: rembg/sessions/sam.py
================================================
import os
from copy import deepcopy
from typing import List
import numpy as np
import onnxruntime as ort
import pooch
from jsonschema import validate
from PIL import Image
from PIL.Image import Image as PILImage
from scipy.ndimage import map_coordinates
from .base import BaseSession
def warp_affine(
image: np.ndarray, matrix: np.ndarray, output_shape: tuple
) -> np.ndarray:
"""
Apply affine transformation to an image (matching cv2.warpAffine behavior).
cv2.warpAffine maps source coordinates to destination coordinates:
dst(M @ [x, y, 1]^T) = src(x, y)
So to fill dst(x', y'), we compute the inverse:
src_coords = M^(-1) @ [x', y', 1]^T
Args:
image: Input image (H, W) or (H, W, C)
matrix: 2x3 affine transformation matrix
output_shape: (height, width) of output
Returns:
Transformed image
"""
h, w = output_shape
# Build full 3x3 matrix and compute inverse
M_full = np.vstack([matrix, [0, 0, 1]])
M_inv = np.linalg.inv(M_full)[:2]
# Create output coordinate grid
cols = np.arange(w)
rows = np.arange(h)
x_coords, y_coords = np.meshgrid(cols, rows)
# Apply inverse transform to get source coordinates
src_x = M_inv[0, 0] * x_coords + M_inv[0, 1] * y_coords + M_inv[0, 2]
src_y = M_inv[1, 0] * x_coords + M_inv[1, 1] * y_coords + M_inv[1, 2]
if image.ndim == 2:
result = map_coordinates(
image.astype(np.float64), [src_y, src_x], order=1, mode="constant", cval=0
)
else:
result = np.zeros((h, w, image.shape[2]), dtype=np.float64)
for c in range(image.shape[2]):
result[:, :, c] = map_coordinates(
image[:, :, c].astype(np.float64),
[src_y, src_x],
order=1,
mode="constant",
cval=0,
)
return result.astype(image.dtype)
def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int):
scale = long_side_length * 1.0 / max(oldh, oldw)
newh, neww = oldh * scale, oldw * scale
neww = int(neww + 0.5)
newh = int(newh + 0.5)
return (newh, neww)
def apply_coords(coords: np.ndarray, original_size, target_length):
old_h, old_w = original_size
new_h, new_w = get_preprocess_shape(
original_size[0], original_size[1], target_length
)
coords = deepcopy(coords).astype(float)
coords[..., 0] = coords[..., 0] * (new_w / old_w)
coords[..., 1] = coords[..., 1] * (new_h / old_h)
return coords
def get_input_points(prompt):
points = []
labels = []
for mark in prompt:
if mark["type"] == "point":
points.append(mark["data"])
labels.append(mark["label"])
elif mark["type"] == "rectangle":
points.append([mark["data"][0], mark["data"][1]])
points.append([mark["data"][2], mark["data"][3]])
labels.append(2)
labels.append(3)
points, labels = np.array(points), np.array(labels)
return points, labels
def transform_masks(masks, original_size, transform_matrix):
output_masks = []
for batch in range(masks.shape[0]):
batch_masks = []
for mask_id in range(masks.shape[1]):
mask = masks[batch, mask_id]
mask = warp_affine(
mask,
transform_matrix[:2],
(original_size[0], original_size[1]),
)
batch_masks.append(mask)
output_masks.append(batch_masks)
return np.array(output_masks)
class SamSession(BaseSession):
"""
This class represents a session for the Sam model.
Args:
model_name (str): The name of the model.
sess_opts (ort.SessionOptions): The session options.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
"""
def __init__(
self,
model_name: str,
sess_opts: ort.SessionOptions,
*args,
**kwargs,
):
"""
Initialize a new SamSession with the given model name and session options.
Args:
model_name (str): The name of the model.
sess_opts (ort.SessionOptions): The session options.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
"""
self.model_name = model_name
paths = self.__class__.download_models(*args, **kwargs)
self.encoder = ort.InferenceSession(
str(paths[0]),
sess_options=sess_opts,
)
self.decoder = ort.InferenceSession(
str(paths[1]),
sess_options=sess_opts,
)
def predict(
self,
img: PILImage,
*args,
**kwargs,
) -> List[PILImage]:
"""
Predict masks for an input image.
This function takes an image as input and performs various preprocessing steps on the image. It then runs the image through an encoder to obtain an image embedding. The function also takes input labels and points as additional arguments. It concatenates the input points and labels with padding and transforms them. It creates an empty mask input and an indicator for no mask. The function then passes the image embedding, point coordinates, point labels, mask input, and has mask input to a decoder. The decoder generates masks based on the input and returns them as a list of images.
Parameters:
img (PILImage): The input image.
*args: Additional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: A list of masks generated by the decoder.
"""
prompt = kwargs.get(
"sam_prompt",
[
{
"type": "point",
"label": 1,
"data": [int(img.width / 2), int(img.height / 2)],
}
],
)
schema = {
"type": "array",
"items": {
"type": "object",
"properties": {
"type": {"type": "string"},
"label": {"type": "integer"},
"data": {
"type": "array",
"items": {"type": "number"},
},
},
},
}
validate(instance=prompt, schema=schema)
target_size = 1024
input_size = (684, 1024)
encoder_input_name = self.encoder.get_inputs()[0].name
img = img.convert("RGB")
cv_image = np.array(img)
original_size = cv_image.shape[:2]
scale_x = input_size[1] / cv_image.shape[1]
scale_y = input_size[0] / cv_image.shape[0]
scale = min(scale_x, scale_y)
transform_matrix = np.array(
[
[scale, 0, 0],
[0, scale, 0],
[0, 0, 1],
]
)
cv_image = warp_affine(
cv_image,
transform_matrix[:2],
(input_size[0], input_size[1]),
)
## encoder
encoder_inputs = {
encoder_input_name: cv_image.astype(np.float32),
}
encoder_output = self.encoder.run(None, encoder_inputs)
image_embedding = encoder_output[0]
embedding = {
"image_embedding": image_embedding,
"original_size": original_size,
"transform_matrix": transform_matrix,
}
## decoder
input_points, input_labels = get_input_points(prompt)
onnx_coord = np.concatenate([input_points, np.array([[0.0, 0.0]])], axis=0)[
None, :, :
]
onnx_label = np.concatenate([input_labels, np.array([-1])], axis=0)[
None, :
].astype(np.float32)
onnx_coord = apply_coords(onnx_coord, input_size, target_size).astype(
np.float32
)
onnx_coord = np.concatenate(
[
onnx_coord,
np.ones((1, onnx_coord.shape[1], 1), dtype=np.float32),
],
axis=2,
)
onnx_coord = np.matmul(onnx_coord, transform_matrix.T)
onnx_coord = onnx_coord[:, :, :2].astype(np.float32)
onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
onnx_has_mask_input = np.zeros(1, dtype=np.float32)
decoder_inputs = {
"image_embeddings": image_embedding,
"point_coords": onnx_coord,
"point_labels": onnx_label,
"mask_input": onnx_mask_input,
"has_mask_input": onnx_has_mask_input,
"orig_im_size": np.array(input_size, dtype=np.float32),
}
masks, _, _ = self.decoder.run(None, decoder_inputs)
inv_transform_matrix = np.linalg.inv(transform_matrix)
masks = transform_masks(masks, original_size, inv_transform_matrix)
mask = np.zeros((masks.shape[2], masks.shape[3], 3), dtype=np.uint8)
for m in masks[0, :, :, :]:
mask[m > 0.0] = [255, 255, 255]
return [Image.fromarray(mask).convert("L")]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Class method to download ONNX model files.
This method is responsible for downloading two ONNX model files from specified URLs and saving them locally. The downloaded files are saved with the naming convention 'name_encoder.onnx' and 'name_decoder.onnx', where 'name' is the value returned by the 'name' method.
Parameters:
cls: The class object.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
tuple: A tuple containing the file paths of the downloaded encoder and decoder models.
"""
model_name = kwargs.get("sam_model", "sam_vit_b_01ec64")
quant = kwargs.get("sam_quant", False)
fname_encoder = f"{model_name}.encoder.onnx"
fname_decoder = f"{model_name}.decoder.onnx"
if quant:
fname_encoder = f"{model_name}.encoder.quant.onnx"
fname_decoder = f"{model_name}.decoder.quant.onnx"
pooch.retrieve(
f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/{fname_encoder}",
None,
fname=fname_encoder,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
pooch.retrieve(
f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/{fname_decoder}",
None,
fname=fname_decoder,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
if fname_encoder == "sam_vit_h_4b8939.encoder.onnx" and not os.path.exists(
os.path.join(
cls.u2net_home(*args, **kwargs), "sam_vit_h_4b8939.encoder_data.bin"
)
):
content = bytearray()
for i in range(1, 4):
pooch.retrieve(
f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/sam_vit_h_4b8939.encoder_data.{i}.bin",
None,
fname=f"sam_vit_h_4b8939.encoder_data.{i}.bin",
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
fbin = os.path.join(
cls.u2net_home(*args, **kwargs),
f"sam_vit_h_4b8939.encoder_data.{i}.bin",
)
content.extend(open(fbin, "rb").read())
os.remove(fbin)
with open(
os.path.join(
cls.u2net_home(*args, **kwargs),
"sam_vit_h_4b8939.encoder_data.bin",
),
"wb",
) as fp:
fp.write(content)
return (
os.path.join(cls.u2net_home(*args, **kwargs), fname_encoder),
os.path.join(cls.u2net_home(*args, **kwargs), fname_decoder),
)
@classmethod
def name(cls, *args, **kwargs):
"""
Class method to return a string value.
This method returns the string value 'sam'.
Parameters:
cls: The class object.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The string value 'sam'.
"""
return "sam"
================================================
FILE: rembg/sessions/silueta.py
================================================
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 SiluetaSession(BaseSession):
"""This is a class representing a SiluetaSession object."""
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predict the mask of the input image.
This method takes an image as input, preprocesses it, and performs a prediction to generate a mask. The generated mask is then post-processed and returned as a list of PILImage objects.
Parameters:
img (PILImage): The input image to be processed.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
List[PILImage]: A list of post-processed masks.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
),
)
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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Download the pre-trained model file.
This method downloads the pre-trained model file from a specified URL. The file is saved to the U2NET home directory.
Parameters:
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name()}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:55e59e0d8062d2f5d013f4725ee84782"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Return the name of the model.
This method returns the name of the Silueta model.
Parameters:
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The name of the model.
"""
return "silueta"
================================================
FILE: rembg/sessions/u2net.py
================================================
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 U2netSession(BaseSession):
"""
This class represents a U2net session, which is a subclass of BaseSession.
"""
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predicts the output masks for the input image using the inner session.
Parameters:
img (PILImage): The input image.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: The list of output masks.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
),
)
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.clip(0, 1) * 255).astype("uint8"), mode="L")
mask = mask.resize(img.size, Image.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the U2net model file from a specific URL and saves it.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The path to the downloaded model file.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:60024c5c889badc19c04ad937298a77b"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the U2net session.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the session.
"""
return "u2net"
================================================
FILE: rembg/sessions/u2net_cloth_seg.py
================================================
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
palette1 = [
0,
0,
0,
255,
255,
255,
0,
0,
0,
0,
0,
0,
]
palette2 = [
0,
0,
0,
0,
0,
0,
255,
255,
255,
0,
0,
0,
]
palette3 = [
0,
0,
0,
0,
0,
0,
0,
0,
0,
255,
255,
255,
]
class Unet2ClothSession(BaseSession):
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predict the cloth category of an image.
This method takes an image as input and predicts the cloth category of the image.
The method uses the inner_session to make predictions using a pre-trained model.
The predicted mask is then converted to an image and resized to match the size of the input image.
Depending on the cloth category specified in the method arguments, the method applies different color palettes to the mask and appends the resulting images to a list.
Parameters:
img (PILImage): The input image.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: A list of images representing the predicted masks.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (768, 768)
),
)
pred = np.argmax(ort_outs[0], axis=1, keepdims=True)
pred = np.squeeze(pred, 0)
pred = np.squeeze(pred, 0)
mask = Image.fromarray(pred.astype("uint8"), mode="L")
mask = mask.resize(img.size, Image.Resampling.LANCZOS)
masks = []
cloth_category = kwargs.get("cc") or kwargs.get("cloth_category")
def upper_cloth():
mask1 = mask.copy()
mask1.putpalette(palette1)
mask1 = mask1.convert("RGB").convert("L")
masks.append(mask1)
def lower_cloth():
mask2 = mask.copy()
mask2.putpalette(palette2)
mask2 = mask2.convert("RGB").convert("L")
masks.append(mask2)
def full_cloth():
mask3 = mask.copy()
mask3.putpalette(palette3)
mask3 = mask3.convert("RGB").convert("L")
masks.append(mask3)
if cloth_category == "upper":
upper_cloth()
elif cloth_category == "lower":
lower_cloth()
elif cloth_category == "full":
full_cloth()
else:
upper_cloth()
lower_cloth()
full_cloth()
return masks
@classmethod
def download_models(cls, *args, **kwargs):
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:2434d1f3cb744e0e49386c906e5a08bb"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
return "u2net_cloth_seg"
================================================
FILE: rembg/sessions/u2net_custom.py
================================================
import os
from typing import List
import numpy as np
import onnxruntime as ort
import pooch
from PIL import Image
from PIL.Image import Image as PILImage
from .base import BaseSession
class U2netCustomSession(BaseSession):
"""This is a class representing a custom session for the U2net model."""
def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *args, **kwargs):
"""
Initialize a new U2netCustomSession object.
Parameters:
model_name (str): The name of the model.
sess_opts (ort.SessionOptions): The session options.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Raises:
ValueError: If model_path is None.
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path is required")
super().__init__(model_name, sess_opts, *args, **kwargs)
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predict the segmentation mask for the input image.
Parameters:
img (PILImage): The input image.
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: A list of PILImage objects representing the segmentation mask.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
),
)
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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Download the model files.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The absolute path to the model files.
"""
model_path = kwargs.get("model_path")
if model_path is None:
raise ValueError("model_path is required")
return os.path.abspath(os.path.expanduser(model_path))
@classmethod
def name(cls, *args, **kwargs):
"""
Get the name of the model.
Parameters:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Returns:
str: The name of the model.
"""
return "u2net_custom"
================================================
FILE: rembg/sessions/u2net_human_seg.py
================================================
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 U2netHumanSegSession(BaseSession):
"""
This class represents a session for performing human segmentation using the U2Net model.
"""
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predicts human segmentation masks for the input image.
Parameters:
img (PILImage): The input image.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
List[PILImage]: A list of predicted masks.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
),
)
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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the U2Net model weights.
Parameters:
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The path to the downloaded model weights.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:c09ddc2e0104f800e3e1bb4652583d1f"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the U2Net model.
Parameters:
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The name of the model.
"""
return "u2net_human_seg"
================================================
FILE: rembg/sessions/u2netp.py
================================================
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 U2netpSession(BaseSession):
"""This class represents a session for using the U2netp model."""
def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
"""
Predicts the mask for the given image using the U2netp model.
Parameters:
img (PILImage): The input image.
Returns:
List[PILImage]: The predicted mask.
"""
ort_outs = self.inner_session.run(
None,
self.normalize(
img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
),
)
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.Resampling.LANCZOS)
return [mask]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Downloads the U2netp model.
Returns:
str: The path to the downloaded model.
"""
fname = f"{cls.name(*args, **kwargs)}.onnx"
pooch.retrieve(
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx",
(
None
if cls.checksum_disabled(*args, **kwargs)
else "md5:8e83ca70e441ab06c318d82300c84806"
),
fname=fname,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
return os.path.join(cls.u2net_home(*args, **kwargs), fname)
@classmethod
def name(cls, *args, **kwargs):
"""
Returns the name of the U2netp model.
Returns:
str: The name of the model.
"""
return "u2netp"
================================================
FILE: rembg.ipynb
================================================
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": [],
"gpuType": "T4"
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "hF9llNyHkiRB",
"outputId": "bd4e1cc0-f871-4c3f-d6e3-b503fe170f71"
},
"outputs": [
],
"source": [
"! pip install \"rembg[gpu,cli]\"\n",
"! git clone https://huggingface.co/spaces/KenjieDec/RemBG\n",
"%cd RemBG\n",
"!python app.py"
]
}
]
}
================================================
FILE: rembg.py
================================================
from rembg.cli import main
if __name__ == "__main__":
main()
================================================
FILE: rembg.spec
================================================
# -*- mode: python ; coding: utf-8 -*-
from PyInstaller.utils.hooks import collect_data_files, collect_dynamic_libs
datas = []
datas += collect_data_files('gradio_client')
datas += collect_data_files('gradio')
datas += collect_data_files('safehttpx')
datas += collect_data_files('groovy')
binaries = []
# Collect onnxruntime (works for both CPU and GPU versions)
# The pip packages are named differently (onnxruntime vs onnxruntime-gpu)
# but both install the Python module as 'onnxruntime'
try:
datas += collect_data_files('onnxruntime')
binaries += collect_dynamic_libs('onnxruntime')
except Exception:
pass
a = Analysis(
['rembg.py'],
pathex=[],
binaries=binaries,
datas=datas,
hiddenimports=[
# Core dependencies
'numpy',
'PIL',
'scipy',
'scipy.ndimage',
'skimage',
'skimage.morphology',
'pymatting',
'pymatting.alpha',
'pymatting.foreground',
'pymatting.util',
'tqdm',
'pooch',
'jsonschema',
'onnxruntime',
# CLI dependencies
'click',
'uvicorn',
'fastapi',
'starlette',
'starlette.responses',
'aiohttp',
'asyncer',
'filetype',
'gradio',
'watchdog',
'sniffio',
'multipart',
],
hookspath=[],
hooksconfig={},
runtime_hooks=[],
excludes=[],
noarchive=False,
module_collection_mode={
'gradio': 'py',
},
)
pyz = PYZ(a.pure)
exe = EXE(
pyz,
a.scripts,
[],
exclude_binaries=True,
name='rembg',
debug=False,
bootloader_ignore_signals=False,
strip=False,
upx=True,
console=True,
disable_windowed_traceback=False,
argv_emulation=False,
target_arch=None,
codesign_identity=None,
entitlements_file=None,
)
coll = COLLECT(
exe,
a.binaries,
a.datas,
strip=False,
upx=True,
upx_exclude=[],
name='rembg',
)
================================================
FILE: tests/test_remove.py
================================================
from io import BytesIO
from pathlib import Path
from imagehash import phash as hash_img
from PIL import Image
from rembg import new_session, remove
here = Path(__file__).parent.resolve()
failures_dir = here / "failures"
failures_dir.mkdir(exist_ok=True)
def test_remove():
kwargs = {
"sam": {
"anime-girl-1" : {
"sam_prompt" :[{"type": "point", "data": [400, 165], "label": 1}],
}
}
}
for model in [
"u2net",
"u2netp",
"u2net_human_seg",
"u2net_cloth_seg",
"silueta",
"isnet-general-use",
"isnet-anime",
"sam",
"birefnet-general",
"birefnet-general-lite",
"birefnet-portrait",
"birefnet-dis",
"birefnet-hrsod",
"birefnet-cod",
"birefnet-massive"
]:
for picture in ["anime-girl-1"]:
image_path = Path(here / "fixtures" / f"{picture}.jpg")
image = image_path.read_bytes()
actual = remove(image, session=new_session(model), **kwargs.get(model, {}).get(picture, {}))
actual_hash = hash_img(Image.open(BytesIO(actual)))
expected_path = Path(here / "results" / f"{picture}.{model}.png")
# Uncomment to update the expected results
# f = open(expected_path, "wb")
# f.write(actual)
# f.close()
expected = expected_path.read_bytes()
expected_hash = hash_img(Image.open(BytesIO(expected)))
print(f"image_path: {image_path}")
print(f"expected_path: {expected_path}")
print(f"actual_hash: {actual_hash}")
print(f"expected_hash: {expected_hash}")
print(f"actual_hash == expected_hash: {actual_hash == expected_hash}")
print("---\n")
if actual_hash != expected_hash:
# Salva as imagens que falharam para comparação
actual_failure_path = failures_dir / f"{picture}.{model}.actual.png"
expected_failure_path = failures_dir / f"{picture}.{model}.expected.png"
with open(actual_failure_path, "wb") as f:
f.write(actual)
with open(expected_failure_path, "wb") as f:
f.write(expected)
print(f"FAILURE: Saved comparison images to {failures_dir}")
assert actual_hash == expected_hash
gitextract_xxyh1rmd/
├── .dockerignore
├── .editorconfig
├── .gitattributes
├── .github/
│ ├── FUNDING.yml
│ ├── ISSUE_TEMPLATE/
│ │ ├── bug_report.md
│ │ └── feature_request.md
│ └── workflows/
│ ├── close_inactive_issues.yml
│ ├── lint_python.yml
│ ├── publish_docker.yml
│ ├── publish_pypi.yml
│ └── windows_installer.yml
├── .gitignore
├── .markdownlint.yaml
├── .python-version
├── CITATION.cff
├── Dockerfile
├── Dockerfile_nvidia_cuda_cudnn_gpu
├── LICENSE.txt
├── MANIFEST.in
├── README.md
├── USAGE.md
├── _build-exe.ps1
├── _modpath.iss
├── _setup-cpu.iss
├── _setup-gpu.iss
├── docker-compose.yml
├── man/
│ └── rembg.1
├── pyproject.toml
├── pytest.ini
├── rembg/
│ ├── __init__.py
│ ├── bg.py
│ ├── cli.py
│ ├── commands/
│ │ ├── __init__.py
│ │ ├── b_command.py
│ │ ├── d_command.py
│ │ ├── i_command.py
│ │ ├── p_command.py
│ │ └── s_command.py
│ ├── session_factory.py
│ └── sessions/
│ ├── __init__.py
│ ├── base.py
│ ├── ben_custom.py
│ ├── birefnet_cod.py
│ ├── birefnet_dis.py
│ ├── birefnet_general.py
│ ├── birefnet_general_lite.py
│ ├── birefnet_hrsod.py
│ ├── birefnet_massive.py
│ ├── birefnet_portrait.py
│ ├── bria_rmbg.py
│ ├── dis_anime.py
│ ├── dis_custom.py
│ ├── dis_general_use.py
│ ├── sam.py
│ ├── silueta.py
│ ├── u2net.py
│ ├── u2net_cloth_seg.py
│ ├── u2net_custom.py
│ ├── u2net_human_seg.py
│ └── u2netp.py
├── rembg.ipynb
├── rembg.py
├── rembg.spec
└── tests/
└── test_remove.py
SYMBOL INDEX (107 symbols across 29 files)
FILE: rembg/bg.py
class ReturnType (line 39) | class ReturnType(Enum):
function alpha_matting_cutout (line 45) | def alpha_matting_cutout(
function naive_cutout (line 98) | def naive_cutout(img: PILImage, mask: PILImage) -> PILImage:
function putalpha_cutout (line 114) | def putalpha_cutout(img: PILImage, mask: PILImage) -> PILImage:
function get_concat_v_multi (line 129) | def get_concat_v_multi(imgs: List[PILImage]) -> PILImage:
function get_concat_v (line 145) | def get_concat_v(img1: PILImage, img2: PILImage) -> PILImage:
function post_process (line 162) | def post_process(mask: np.ndarray) -> np.ndarray:
function apply_background_color (line 175) | def apply_background_color(img: PILImage, color: Tuple[int, int, int, in...
function fix_image_orientation (line 192) | def fix_image_orientation(img: PILImage) -> PILImage:
function download_models (line 205) | def download_models(models: tuple[str, ...]) -> None:
function remove (line 226) | def remove(
FILE: rembg/cli.py
function main (line 30) | def main() -> None:
FILE: rembg/commands/b_command.py
function b_command (line 97) | def b_command(
FILE: rembg/commands/d_command.py
function d_command (line 11) | def d_command(models: tuple[str, ...]) -> None:
FILE: rembg/commands/i_command.py
function i_command (line 87) | def i_command(model: str, extras: str, input: IO, output: IO, **kwargs) ...
FILE: rembg/commands/p_command.py
function p_command (line 120) | def p_command(
FILE: rembg/commands/s_command.py
function s_command (line 58) | def s_command(port: int, host: str, log_level: str, threads: int) -> None:
FILE: rembg/session_factory.py
function new_session (line 11) | def new_session(model_name: str = "u2net", *args, **kwargs) -> BaseSession:
FILE: rembg/sessions/base.py
class BaseSession (line 10) | class BaseSession:
method __init__ (line 13) | def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *ar...
method normalize (line 40) | def normalize(
method predict (line 67) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method checksum_disabled (line 71) | def checksum_disabled(cls, *args, **kwargs):
method u2net_home (line 75) | def u2net_home(cls, *args, **kwargs):
method download_models (line 83) | def download_models(cls, *args, **kwargs):
method name (line 87) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/ben_custom.py
class BenCustomSession (line 12) | class BenCustomSession(BaseSession):
method __init__ (line 15) | def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *ar...
method predict (line 31) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 63) | def download_models(cls, *args, **kwargs):
method name (line 81) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/birefnet_cod.py
class BiRefNetSessionCOD (line 8) | class BiRefNetSessionCOD(BiRefNetSessionGeneral):
method download_models (line 14) | def download_models(cls, *args, **kwargs):
method name (line 41) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/birefnet_dis.py
class BiRefNetSessionDIS (line 8) | class BiRefNetSessionDIS(BiRefNetSessionGeneral):
method download_models (line 14) | def download_models(cls, *args, **kwargs):
method name (line 41) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/birefnet_general.py
class BiRefNetSessionGeneral (line 12) | class BiRefNetSessionGeneral(BaseSession):
method sigmoid (line 17) | def sigmoid(self, mat):
method predict (line 20) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 53) | def download_models(cls, *args, **kwargs):
method name (line 80) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/birefnet_general_lite.py
class BiRefNetSessionGeneralLite (line 8) | class BiRefNetSessionGeneralLite(BiRefNetSessionGeneral):
method download_models (line 14) | def download_models(cls, *args, **kwargs):
method name (line 41) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/birefnet_hrsod.py
class BiRefNetSessionHRSOD (line 8) | class BiRefNetSessionHRSOD(BiRefNetSessionGeneral):
method download_models (line 14) | def download_models(cls, *args, **kwargs):
method name (line 41) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/birefnet_massive.py
class BiRefNetSessionMassive (line 8) | class BiRefNetSessionMassive(BiRefNetSessionGeneral):
method download_models (line 14) | def download_models(cls, *args, **kwargs):
method name (line 41) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/birefnet_portrait.py
class BiRefNetSessionPortrait (line 8) | class BiRefNetSessionPortrait(BiRefNetSessionGeneral):
method download_models (line 14) | def download_models(cls, *args, **kwargs):
method name (line 41) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/bria_rmbg.py
class BriaRmBgSession (line 12) | class BriaRmBgSession(BaseSession):
method predict (line 17) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 50) | def download_models(cls, *args, **kwargs):
method name (line 77) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/dis_anime.py
class DisSession (line 12) | class DisSession(BaseSession):
method predict (line 17) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 48) | def download_models(cls, *args, **kwargs):
method name (line 75) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/dis_custom.py
class DisCustomSession (line 12) | class DisCustomSession(BaseSession):
method __init__ (line 15) | def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *ar...
method predict (line 31) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 62) | def download_models(cls, *args, **kwargs):
method name (line 80) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/dis_general_use.py
class DisSession (line 12) | class DisSession(BaseSession):
method predict (line 13) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 44) | def download_models(cls, *args, **kwargs):
method name (line 73) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/sam.py
function warp_affine (line 16) | def warp_affine(
function get_preprocess_shape (line 69) | def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int):
function apply_coords (line 78) | def apply_coords(coords: np.ndarray, original_size, target_length):
function get_input_points (line 91) | def get_input_points(prompt):
function transform_masks (line 109) | def transform_masks(masks, original_size, transform_matrix):
class SamSession (line 127) | class SamSession(BaseSession):
method __init__ (line 138) | def __init__(
method predict (line 166) | def predict(
method download_models (line 299) | def download_models(cls, *args, **kwargs):
method name (line 377) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/silueta.py
class SiluetaSession (line 12) | class SiluetaSession(BaseSession):
method predict (line 15) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 50) | def download_models(cls, *args, **kwargs):
method name (line 79) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/u2net.py
class U2netSession (line 12) | class U2netSession(BaseSession):
method predict (line 17) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 50) | def download_models(cls, *args, **kwargs):
method name (line 77) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/u2net_cloth_seg.py
class Unet2ClothSession (line 57) | class Unet2ClothSession(BaseSession):
method predict (line 58) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 125) | def download_models(cls, *args, **kwargs):
method name (line 142) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/u2net_custom.py
class U2netCustomSession (line 13) | class U2netCustomSession(BaseSession):
method __init__ (line 16) | def __init__(self, model_name: str, sess_opts: ort.SessionOptions, *ar...
method predict (line 35) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 68) | def download_models(cls, *args, **kwargs):
method name (line 86) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/u2net_human_seg.py
class U2netHumanSegSession (line 12) | class U2netHumanSegSession(BaseSession):
method predict (line 17) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 50) | def download_models(cls, *args, **kwargs):
method name (line 77) | def name(cls, *args, **kwargs):
FILE: rembg/sessions/u2netp.py
class U2netpSession (line 12) | class U2netpSession(BaseSession):
method predict (line 15) | def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
method download_models (line 46) | def download_models(cls, *args, **kwargs):
method name (line 69) | def name(cls, *args, **kwargs):
FILE: tests/test_remove.py
function test_remove (line 13) | def test_remove():
Condensed preview — 64 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (156K chars).
[
{
"path": ".dockerignore",
"chars": 60,
"preview": "*\n!rembg\n!pyproject.toml\n!poetry.lock\n!README.md\n!.git\n.env\n"
},
{
"path": ".editorconfig",
"chars": 176,
"preview": "# https://editorconfig.org/\n\nroot = true\n\n[*]\nindent_style = space\nindent_size = 4\ninsert_final_newline = true\ntrim_trai"
},
{
"path": ".gitattributes",
"chars": 31,
"preview": "rembg/_version.py export-subst\n"
},
{
"path": ".github/FUNDING.yml",
"chars": 75,
"preview": "github: [danielgatis]\ncustom: [\"https://www.buymeacoffee.com/danielgatis\"]\n"
},
{
"path": ".github/ISSUE_TEMPLATE/bug_report.md",
"chars": 558,
"preview": "---\nname: Bug report\nabout: Create a report to help us improve\ntitle: \"[BUG] ...\"\nlabels: bug\nassignees: \"\"\n---\n\n**Descr"
},
{
"path": ".github/ISSUE_TEMPLATE/feature_request.md",
"chars": 616,
"preview": "---\nname: Feature request\nabout: Suggest an idea for this project\ntitle: \"[FEATURE] ...\"\nlabels: enhancement\nassignees: "
},
{
"path": ".github/workflows/close_inactive_issues.yml",
"chars": 823,
"preview": "name: Close inactive issues\n\non:\n schedule:\n - cron: \"30 1 * * *\"\n\njobs:\n close_inactive_issues:\n ru"
},
{
"path": ".github/workflows/lint_python.yml",
"chars": 1080,
"preview": "name: Lint\n\non:\n push:\n branches:\n - \"**\"\n pull_request:\n\njobs:\n lint_python:\n runs-on: ubuntu-lates"
},
{
"path": ".github/workflows/publish_docker.yml",
"chars": 1648,
"preview": "name: Publish Docker image\n\non:\n push:\n tags:\n - \"v*.*.*\"\n\njobs:\n publish_docker:\n name: Push Docker image "
},
{
"path": ".github/workflows/publish_pypi.yml",
"chars": 1939,
"preview": "name: Publish to Pypi\n\non:\n push:\n tags:\n - \"v*.*.*\"\n\njobs:\n publish_pypi:\n runs-on: ubun"
},
{
"path": ".github/workflows/windows_installer.yml",
"chars": 1279,
"preview": "name: Build Windows Installer\n\non:\n push:\n tags:\n - \"v*.*.*\"\njobs:\n windows_installer:\n name: B"
},
{
"path": ".gitignore",
"chars": 310,
"preview": "# general things to ignore\nbuild/\ndist/\n.venv/\n.direnv/\n*.egg-info/\n*.egg\n*.py[cod]\n__pycache__/\n*.so\n*~≈\n.env\n.envrc\n.i"
},
{
"path": ".markdownlint.yaml",
"chars": 75,
"preview": "---\ndefault: true\nMD013: false # line-length\nMD033: false # no-inline-html\n"
},
{
"path": ".python-version",
"chars": 7,
"preview": "3.13.9\n"
},
{
"path": "CITATION.cff",
"chars": 523,
"preview": "cff-version: 1.2.0\ntitle: rembg\nmessage: Rembg is a tool to remove images background\ntype: software\nauthors:\n - given-n"
},
{
"path": "Dockerfile",
"chars": 401,
"preview": "FROM python:3.11-slim\n\nWORKDIR /rembg\n\nRUN pip install --upgrade pip && \\\n pip install poetry poetry-dynamic-versioni"
},
{
"path": "Dockerfile_nvidia_cuda_cudnn_gpu",
"chars": 354,
"preview": "FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04\n\nWORKDIR /rembg\n\nRUN apt-get update && apt-get install -y --no-install-r"
},
{
"path": "LICENSE.txt",
"chars": 1069,
"preview": "MIT License\n\nCopyright (c) 2020 Daniel Gatis\n\nPermission is hereby granted, free of charge, to any person obtaining a co"
},
{
"path": "MANIFEST.in",
"chars": 61,
"preview": "include LICENSE.txt\ninclude README.md\ninclude pyproject.toml\n"
},
{
"path": "README.md",
"chars": 15072,
"preview": "<p align=\"center\">\n <img src=\"logo.png\" alt=\"Rembg Logo\" width=\"600\" />\n</p>\n\n<div align=\"center\">\n <p align=\"center\">"
},
{
"path": "USAGE.md",
"chars": 2319,
"preview": "# How to use the remove function\n\n## Load the Image\n\n```python\nfrom PIL import Image\nfrom rembg import new_session, remo"
},
{
"path": "_build-exe.ps1",
"chars": 747,
"preview": "# Install Poetry if not already installed\nif (-not (Get-Command poetry -ErrorAction SilentlyContinue)) {\n pip install"
},
{
"path": "_modpath.iss",
"chars": 6731,
"preview": "// ----------------------------------------------------------------------------\n//\n// Inno Setup Ver:\t5.4.2\n// Script Ve"
},
{
"path": "_setup-cpu.iss",
"chars": 1275,
"preview": "#define MyAppName \"Rembg CPU\"\n#define MyAppVersion \"STABLE\"\n#define MyAppPublisher \"danielgatis\"\n#define MyAppURL \"https"
},
{
"path": "_setup-gpu.iss",
"chars": 1275,
"preview": "#define MyAppName \"Rembg GPU\"\n#define MyAppVersion \"STABLE\"\n#define MyAppPublisher \"danielgatis\"\n#define MyAppURL \"https"
},
{
"path": "docker-compose.yml",
"chars": 258,
"preview": "---\n# You can set variables in .env file in root folder\n#\n# PUBLIC_PORT=7000:7000\n# REPLICAS_COUNT=1\n\nservices:\n app:\n "
},
{
"path": "man/rembg.1",
"chars": 456,
"preview": ".TH REMBG 1 \"Januar 2026\" \"2.0.72\" \"User Commands\"\n.SH NAME\nrembg \\- tool to remove background from images\n.SH SYNOPSIS\n"
},
{
"path": "pyproject.toml",
"chars": 3384,
"preview": "[tool.poetry]\nname = \"rembg\"\nversion = \"0.0.0\" # Managed by poetry-dynamic-versioning\ndescription = \"Remove image backg"
},
{
"path": "pytest.ini",
"chars": 57,
"preview": "[pytest]\nfilterwarnings =\n ignore::DeprecationWarning\n"
},
{
"path": "rembg/__init__.py",
"chars": 359,
"preview": "try:\n from importlib.metadata import PackageNotFoundError, version\n\n try:\n __version__ = version(\"rembg\")\n "
},
{
"path": "rembg/bg.py",
"chars": 11538,
"preview": "import io\nimport sys\nfrom enum import Enum\nfrom typing import Any, List, Optional, Tuple, Union, cast\n\nimport numpy as n"
},
{
"path": "rembg/cli.py",
"chars": 885,
"preview": "import sys\n\n# Fast path for --version (avoid importing heavy dependencies)\nif len(sys.argv) == 2 and sys.argv[1] in (\"--"
},
{
"path": "rembg/commands/__init__.py",
"chars": 370,
"preview": "command_functions = []\n\nfrom .b_command import b_command\nfrom .d_command import d_command\nfrom .i_command import i_comma"
},
{
"path": "rembg/commands/b_command.py",
"chars": 4764,
"preview": "import asyncio\nimport io\nimport json\nimport os\nimport sys\nfrom typing import IO\n\nimport click\nimport PIL\n\nfrom ..bg impo"
},
{
"path": "rembg/commands/d_command.py",
"chars": 273,
"preview": "import click\n\nfrom ..bg import download_models\n\n\n@click.command( # type: ignore\n name=\"d\",\n help=\"download models"
},
{
"path": "rembg/commands/i_command.py",
"chars": 2624,
"preview": "import json\nimport sys\nfrom typing import IO\n\nimport click\n\nfrom ..bg import remove\nfrom ..session_factory import new_se"
},
{
"path": "rembg/commands/p_command.py",
"chars": 5947,
"preview": "import json\nimport pathlib\nimport time\nfrom typing import cast\n\nimport click\nimport filetype\nfrom tqdm import tqdm\nfrom "
},
{
"path": "rembg/commands/s_command.py",
"chars": 10719,
"preview": "import json\nimport os\nimport webbrowser\nfrom typing import Optional, Tuple, cast\n\nimport aiohttp\nimport click\nimport gra"
},
{
"path": "rembg/session_factory.py",
"chars": 1638,
"preview": "import os\nfrom typing import Optional, Type\n\nimport onnxruntime as ort\n\nfrom .sessions import sessions_class\nfrom .sessi"
},
{
"path": "rembg/sessions/__init__.py",
"chars": 2173,
"preview": "from __future__ import annotations\n\nfrom typing import Dict, List\n\nfrom .base import BaseSession\n\nsessions: Dict[str, ty"
},
{
"path": "rembg/sessions/base.py",
"chars": 2858,
"preview": "import os\nfrom typing import Dict, List, Tuple\n\nimport numpy as np\nimport onnxruntime as ort\nfrom PIL import Image\nfrom "
},
{
"path": "rembg/sessions/ben_custom.py",
"chars": 2697,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport onnxruntime as ort\nfrom PIL import Image\nfrom PIL.Image imp"
},
{
"path": "rembg/sessions/birefnet_cod.py",
"chars": 1475,
"preview": "import os\n\nimport pooch\n\nfrom . import BiRefNetSessionGeneral\n\n\nclass BiRefNetSessionCOD(BiRefNetSessionGeneral):\n \"\""
},
{
"path": "rembg/sessions/birefnet_dis.py",
"chars": 1475,
"preview": "import os\n\nimport pooch\n\nfrom . import BiRefNetSessionGeneral\n\n\nclass BiRefNetSessionDIS(BiRefNetSessionGeneral):\n \"\""
},
{
"path": "rembg/sessions/birefnet_general.py",
"chars": 2592,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/birefnet_general_lite.py",
"chars": 1539,
"preview": "import os\n\nimport pooch\n\nfrom . import BiRefNetSessionGeneral\n\n\nclass BiRefNetSessionGeneralLite(BiRefNetSessionGeneral)"
},
{
"path": "rembg/sessions/birefnet_hrsod.py",
"chars": 1491,
"preview": "import os\n\nimport pooch\n\nfrom . import BiRefNetSessionGeneral\n\n\nclass BiRefNetSessionHRSOD(BiRefNetSessionGeneral):\n "
},
{
"path": "rembg/sessions/birefnet_massive.py",
"chars": 1515,
"preview": "import os\n\nimport pooch\n\nfrom . import BiRefNetSessionGeneral\n\n\nclass BiRefNetSessionMassive(BiRefNetSessionGeneral):\n "
},
{
"path": "rembg/sessions/birefnet_portrait.py",
"chars": 1505,
"preview": "import os\n\nimport pooch\n\nfrom . import BiRefNetSessionGeneral\n\n\nclass BiRefNetSessionPortrait(BiRefNetSessionGeneral):\n "
},
{
"path": "rembg/sessions/bria_rmbg.py",
"chars": 2506,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/dis_anime.py",
"chars": 2360,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/dis_custom.py",
"chars": 2696,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport onnxruntime as ort\nfrom PIL import Image\nfrom PIL.Image imp"
},
{
"path": "rembg/sessions/dis_general_use.py",
"chars": 2535,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/sam.py",
"chars": 12534,
"preview": "import os\nfrom copy import deepcopy\nfrom typing import List\n\nimport numpy as np\nimport onnxruntime as ort\nimport pooch\nf"
},
{
"path": "rembg/sessions/silueta.py",
"chars": 2716,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/u2net.py",
"chars": 2445,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/u2net_cloth_seg.py",
"chars": 3478,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/u2net_custom.py",
"chars": 2769,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport onnxruntime as ort\nimport pooch\nfrom PIL import Image\nfrom "
},
{
"path": "rembg/sessions/u2net_human_seg.py",
"chars": 2421,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg/sessions/u2netp.py",
"chars": 1998,
"preview": "import os\nfrom typing import List\n\nimport numpy as np\nimport pooch\nfrom PIL import Image\nfrom PIL.Image import Image as "
},
{
"path": "rembg.ipynb",
"chars": 795,
"preview": "{\n \"nbformat\": 4,\n \"nbformat_minor\": 0,\n \"metadata\": {\n \"colab\": {\n \"provenance\": [],\n \"gpuType\": \"T4\"\n "
},
{
"path": "rembg.py",
"chars": 66,
"preview": "from rembg.cli import main\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "rembg.spec",
"chars": 1984,
"preview": "# -*- mode: python ; coding: utf-8 -*-\nfrom PyInstaller.utils.hooks import collect_data_files, collect_dynamic_libs\n\ndat"
},
{
"path": "tests/test_remove.py",
"chars": 2439,
"preview": "from io import BytesIO\nfrom pathlib import Path\n\nfrom imagehash import phash as hash_img\nfrom PIL import Image\n\nfrom rem"
}
]
About this extraction
This page contains the full source code of the danielgatis/rembg GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 64 files (141.4 KB), approximately 37.2k tokens, and a symbol index with 107 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.