Daniel Gatis 3 年之前
父节点
当前提交
c914e77a7f
共有 3 个文件被更改,包括 3 次插入3 次删除
  1. 0 2
      .github/workflows/lint_python.yml
  2. 1 0
      rembg/bg.py
  3. 2 1
      rembg/detect.py

+ 0 - 2
.github/workflows/lint_python.yml

@@ -8,8 +8,6 @@ jobs:
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v2
-        with:
-            python-version: '3.9'
       - run: pip install --upgrade pip wheel
       - run: pip install bandit black flake8 flake8-bugbear flake8-comprehensions isort safety mypy
       - run: mypy --install-types --non-interactive --ignore-missing-imports ./rembg

+ 1 - 0
rembg/bg.py

@@ -1,4 +1,5 @@
 import io
+from enum import Enum
 from typing import Optional, Union
 
 import numpy as np

+ 2 - 1
rembg/detect.py

@@ -27,8 +27,9 @@ def ort_session(model_name: str) -> ort.InferenceSession:
     home = os.getenv("U2NET_HOME", os.path.join("~", ".u2net"))
     path = Path(home).expanduser() / f"{model_name}.onnx"
     path.parents[0].mkdir(parents=True, exist_ok=True)
+    hashing = hashlib.new("md5", path.read_bytes(), usedforsecurity=False)
 
-    if not (path.exists() and hashlib.md5(path.read_bytes()).hexdigest() == md5):
+    if not (path.exists() and hashing.hexdigest() == md5):
         with redirect_stdout(sys.stderr):
             gdown.download(url, str(path), use_cookies=False)