mirror of
https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg
synced 2025-08-15 10:45:59 +08:00
update deps / remove gdown
This commit is contained in:
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).**
|
**If this project has helped you, please consider making a [donation](https://www.buymeacoffee.com/danielgatis).**
|
||||||
|
|
||||||
|
### Requirements
|
||||||
|
|
||||||
|
```
|
||||||
|
python: >=3.7, <3.11
|
||||||
|
```
|
||||||
|
|
||||||
### Installation
|
### Installation
|
||||||
|
|
||||||
CPU support:
|
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:
|
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.
|
- 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://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.
|
- 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://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_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://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_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
|
#### How to train your own model
|
||||||
|
|
||||||
|
@ -5,7 +5,7 @@ from contextlib import redirect_stdout
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Type
|
from typing import Type
|
||||||
|
|
||||||
import gdown
|
import pooch
|
||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
|
|
||||||
from .session_base import BaseSession
|
from .session_base import BaseSession
|
||||||
@ -18,39 +18,40 @@ def new_session(model_name: str) -> BaseSession:
|
|||||||
|
|
||||||
if model_name == "u2netp":
|
if model_name == "u2netp":
|
||||||
md5 = "8e83ca70e441ab06c318d82300c84806"
|
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
|
session_class = SimpleSession
|
||||||
elif model_name == "u2net":
|
elif model_name == "u2net":
|
||||||
md5 = "60024c5c889badc19c04ad937298a77b"
|
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
|
session_class = SimpleSession
|
||||||
elif model_name == "u2net_human_seg":
|
elif model_name == "u2net_human_seg":
|
||||||
md5 = "c09ddc2e0104f800e3e1bb4652583d1f"
|
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
|
session_class = SimpleSession
|
||||||
elif model_name == "u2net_cloth_seg":
|
elif model_name == "u2net_cloth_seg":
|
||||||
md5 = "2434d1f3cb744e0e49386c906e5a08bb"
|
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
|
session_class = ClothSession
|
||||||
else:
|
else:
|
||||||
assert AssertionError(
|
assert AssertionError(
|
||||||
"Choose between u2net, u2netp, u2net_human_seg or u2net_cloth_seg"
|
"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")
|
"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():
|
fname = f"{model_name}.onnx"
|
||||||
with redirect_stdout(sys.stderr):
|
path = Path(u2net_home).expanduser()
|
||||||
gdown.download(url, str(path), use_cookies=False)
|
full_path = Path(u2net_home).expanduser() / fname
|
||||||
else:
|
|
||||||
hashing = hashlib.new("md5", path.read_bytes(), usedforsecurity=False)
|
pooch.retrieve(
|
||||||
if hashing.hexdigest() != md5:
|
url,
|
||||||
with redirect_stdout(sys.stderr):
|
f"md5:{md5}",
|
||||||
gdown.download(url, str(path), use_cookies=False)
|
fname=fname,
|
||||||
|
path=Path(u2net_home).expanduser(),
|
||||||
|
progressbar=True
|
||||||
|
)
|
||||||
|
|
||||||
sess_opts = ort.SessionOptions()
|
sess_opts = ort.SessionOptions()
|
||||||
|
|
||||||
@ -60,6 +61,6 @@ def new_session(model_name: str) -> BaseSession:
|
|||||||
return session_class(
|
return session_class(
|
||||||
model_name,
|
model_name,
|
||||||
ort.InferenceSession(
|
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
|
aiohttp==3.8.1
|
||||||
asyncer==0.0.1
|
asyncer==0.0.2
|
||||||
click==8.1.3
|
click==8.1.3
|
||||||
fastapi==0.80.0
|
fastapi==0.87.0
|
||||||
filetype==1.1.0
|
filetype==1.2.0
|
||||||
gdown==4.5.1
|
pooch==1.6.0
|
||||||
imagehash==4.2.1
|
imagehash==4.3.1
|
||||||
numpy==1.21.6
|
numpy==1.23.5
|
||||||
onnxruntime==1.12.1
|
onnxruntime==1.13.1
|
||||||
opencv-python-headless==4.6.0.66
|
opencv-python-headless==4.6.0.66
|
||||||
pillow==9.2.0
|
pillow==9.3.0
|
||||||
pymatting==1.1.8
|
pymatting==1.1.8
|
||||||
python-multipart==0.0.5
|
python-multipart==0.0.5
|
||||||
scikit-image==0.19.3
|
scikit-image==0.19.3
|
||||||
scipy==1.7.3
|
scipy==1.9.3
|
||||||
tqdm==4.64.0
|
tqdm==4.64.1
|
||||||
uvicorn==0.18.3
|
uvicorn==0.20.0
|
||||||
watchdog==2.1.9
|
watchdog==2.1.9
|
||||||
|
14
setup.py
14
setup.py
@ -27,10 +27,22 @@ setup(
|
|||||||
author_email="danielgatis@gmail.com",
|
author_email="danielgatis@gmail.com",
|
||||||
classifiers=[
|
classifiers=[
|
||||||
"License :: OSI Approved :: MIT License",
|
"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",
|
keywords="remove, background, u2net",
|
||||||
packages=["rembg"],
|
packages=["rembg"],
|
||||||
python_requires=">=3.7",
|
python_requires=">=3.7, <3.11",
|
||||||
install_requires=requireds,
|
install_requires=requireds,
|
||||||
entry_points={
|
entry_points={
|
||||||
"console_scripts": [
|
"console_scripts": [
|
||||||
|
Loading…
x
Reference in New Issue
Block a user