update deps / remove gdown

This commit is contained in:
Daniel Gatis 2022-11-24 17:24:52 -03:00
parent 8afdfb9dc3
commit 10ce1d0246
4 changed files with 52 additions and 33 deletions

View File

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

View File

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

View File

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

View File

@ -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": [