mirror of
https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg
synced 2025-08-14 19:05:56 +08:00
update deps / remove gdown
This commit is contained in:
Daniel Gatis
parent
8afdfb9dc3
commit
10ce1d0246
14
README.md
14
README.md
@ -37,6 +37,12 @@ Rembg is a tool to remove images background. That is it.
|
||||
|
||||
**If this project has helped you, please consider making a [donation](https://www.buymeacoffee.com/danielgatis).**
|
||||
|
||||
### Requirements
|
||||
|
||||
```
|
||||
python: >=3.7, <3.11
|
||||
```
|
||||
|
||||
### Installation
|
||||
|
||||
CPU support:
|
||||
@ -181,10 +187,10 @@ All models are downloaded and saved in the user home folder in the `.u2net` dire
|
||||
|
||||
The available models are:
|
||||
|
||||
- u2net ([download](https://drive.google.com/uc?id=1tCU5MM1LhRgGou5OpmpjBQbSrYIUoYab) - [alternative](http://depositfiles.com/files/ltxbqa06w), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for general use cases.
|
||||
- u2netp ([download](https://drive.google.com/uc?id=1tNuFmLv0TSNDjYIkjEdeH1IWKQdUA4HR) - [alternative](http://depositfiles.com/files/0y9i0r2fy), [source](https://github.com/xuebinqin/U-2-Net)): A lightweight version of u2net model.
|
||||
- u2net_human_seg ([download](https://drive.google.com/uc?id=1ZfqwVxu-1XWC1xU1GHIP-FM_Knd_AX5j) - [alternative](http://depositfiles.com/files/6spp8qpey), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for human segmentation.
|
||||
- u2net_cloth_seg ([download](https://drive.google.com/uc?id=15rKbQSXQzrKCQurUjZFg8HqzZad8bcyz) - [alternative](http://depositfiles.com/files/l3z3cxetq), [source](https://github.com/levindabhi/cloth-segmentation)): A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body.
|
||||
- u2net ([download](https://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.
|
||||
|
||||
#### How to train your own model
|
||||
|
||||
|
||||
@ -5,7 +5,7 @@ from contextlib import redirect_stdout
|
||||
from pathlib import Path
|
||||
from typing import Type
|
||||
|
||||
import gdown
|
||||
import pooch
|
||||
import onnxruntime as ort
|
||||
|
||||
from .session_base import BaseSession
|
||||
@ -18,39 +18,40 @@ def new_session(model_name: str) -> BaseSession:
|
||||
|
||||
if model_name == "u2netp":
|
||||
md5 = "8e83ca70e441ab06c318d82300c84806"
|
||||
url = "https://drive.google.com/uc?id=1tNuFmLv0TSNDjYIkjEdeH1IWKQdUA4HR"
|
||||
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
|
||||
session_class = SimpleSession
|
||||
elif model_name == "u2net":
|
||||
md5 = "60024c5c889badc19c04ad937298a77b"
|
||||
url = "https://drive.google.com/uc?id=1tCU5MM1LhRgGou5OpmpjBQbSrYIUoYab"
|
||||
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
|
||||
session_class = SimpleSession
|
||||
elif model_name == "u2net_human_seg":
|
||||
md5 = "c09ddc2e0104f800e3e1bb4652583d1f"
|
||||
url = "https://drive.google.com/uc?id=1ZfqwVxu-1XWC1xU1GHIP-FM_Knd_AX5j"
|
||||
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx"
|
||||
session_class = SimpleSession
|
||||
elif model_name == "u2net_cloth_seg":
|
||||
md5 = "2434d1f3cb744e0e49386c906e5a08bb"
|
||||
url = "https://drive.google.com/uc?id=15rKbQSXQzrKCQurUjZFg8HqzZad8bcyz"
|
||||
url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx"
|
||||
session_class = ClothSession
|
||||
else:
|
||||
assert AssertionError(
|
||||
"Choose between u2net, u2netp, u2net_human_seg or u2net_cloth_seg"
|
||||
)
|
||||
|
||||
home = os.getenv(
|
||||
u2net_home = os.getenv(
|
||||
"U2NET_HOME", os.path.join(os.getenv("XDG_DATA_HOME", "~"), ".u2net")
|
||||
)
|
||||
path = Path(home).expanduser() / f"{model_name}.onnx"
|
||||
path.parents[0].mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if not path.exists():
|
||||
with redirect_stdout(sys.stderr):
|
||||
gdown.download(url, str(path), use_cookies=False)
|
||||
else:
|
||||
hashing = hashlib.new("md5", path.read_bytes(), usedforsecurity=False)
|
||||
if hashing.hexdigest() != md5:
|
||||
with redirect_stdout(sys.stderr):
|
||||
gdown.download(url, str(path), use_cookies=False)
|
||||
fname = f"{model_name}.onnx"
|
||||
path = Path(u2net_home).expanduser()
|
||||
full_path = Path(u2net_home).expanduser() / fname
|
||||
|
||||
pooch.retrieve(
|
||||
url,
|
||||
f"md5:{md5}",
|
||||
fname=fname,
|
||||
path=Path(u2net_home).expanduser(),
|
||||
progressbar=True
|
||||
)
|
||||
|
||||
sess_opts = ort.SessionOptions()
|
||||
|
||||
@ -60,6 +61,6 @@ def new_session(model_name: str) -> BaseSession:
|
||||
return session_class(
|
||||
model_name,
|
||||
ort.InferenceSession(
|
||||
str(path), providers=ort.get_available_providers(), sess_options=sess_opts
|
||||
str(full_path), providers=ort.get_available_providers(), sess_options=sess_opts
|
||||
),
|
||||
)
|
||||
|
||||
@ -1,18 +1,18 @@
|
||||
aiohttp==3.8.1
|
||||
asyncer==0.0.1
|
||||
asyncer==0.0.2
|
||||
click==8.1.3
|
||||
fastapi==0.80.0
|
||||
filetype==1.1.0
|
||||
gdown==4.5.1
|
||||
imagehash==4.2.1
|
||||
numpy==1.21.6
|
||||
onnxruntime==1.12.1
|
||||
fastapi==0.87.0
|
||||
filetype==1.2.0
|
||||
pooch==1.6.0
|
||||
imagehash==4.3.1
|
||||
numpy==1.23.5
|
||||
onnxruntime==1.13.1
|
||||
opencv-python-headless==4.6.0.66
|
||||
pillow==9.2.0
|
||||
pillow==9.3.0
|
||||
pymatting==1.1.8
|
||||
python-multipart==0.0.5
|
||||
scikit-image==0.19.3
|
||||
scipy==1.7.3
|
||||
tqdm==4.64.0
|
||||
uvicorn==0.18.3
|
||||
scipy==1.9.3
|
||||
tqdm==4.64.1
|
||||
uvicorn==0.20.0
|
||||
watchdog==2.1.9
|
||||
|
||||
14
setup.py
14
setup.py
@ -27,10 +27,22 @@ setup(
|
||||
author_email="danielgatis@gmail.com",
|
||||
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.7",
|
||||
"Programming Language :: Python :: 3.8",
|
||||
"Programming Language :: Python :: 3.9",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
],
|
||||
keywords="remove, background, u2net",
|
||||
packages=["rembg"],
|
||||
python_requires=">=3.7",
|
||||
python_requires=">=3.7, <3.11",
|
||||
install_requires=requireds,
|
||||
entry_points={
|
||||
"console_scripts": [
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user