diff --git a/README.md b/README.md index 03209a2..b25856e 100644 --- a/README.md +++ b/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 diff --git a/rembg/session_factory.py b/rembg/session_factory.py index 68cab41..a4804b6 100644 --- a/rembg/session_factory.py +++ b/rembg/session_factory.py @@ -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 ), ) diff --git a/requirements.txt b/requirements.txt index 845cd54..f38891a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/setup.py b/setup.py index 4c548f4..f28be8c 100644 --- a/setup.py +++ b/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": [