Răsfoiți Sursa

update deps / remove gdown

Daniel Gatis 2 ani în urmă
părinte
comite
10ce1d0246
4 a modificat fișierele cu 52 adăugiri și 33 ștergeri
  1. 10 4
      README.md
  2. 18 17
      rembg/session_factory.py
  3. 11 11
      requirements.txt
  4. 13 1
      setup.py

+ 10 - 4
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
 

+ 18 - 17
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
         ),
     )

+ 11 - 11
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

+ 13 - 1
setup.py

@@ -27,10 +27,22 @@ setup(
     author_email="[email protected]",
     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": [