Browse Source

fix invalid model fresh install

Daniel Gatis 4 years ago
parent
commit
b07192d16f
3 changed files with 7 additions and 4 deletions
  1. 1 1
      setup.py
  2. 4 3
      src/rembg/cmd/cli.py
  3. 2 0
      src/rembg/cmd/server.py

+ 1 - 1
setup.py

@@ -11,7 +11,7 @@ with open("requirements.txt") as f:
 
 setup(
     name="rembg",
-    version="1.0.25",
+    version="1.0.26",
     description="Remove image background",
     long_description=long_description,
     long_description_content_type="text/markdown",

+ 4 - 3
src/rembg/cmd/cli.py

@@ -15,7 +15,8 @@ def main():
         os.path.expanduser(os.path.join("~", ".u2net")),
     )
     model_choices = [os.path.splitext(os.path.basename(x))[0] for x in set(glob.glob(model_path + "/*"))]
-
+    if len(model_choices) == 0:
+        model_choices = ["u2net", "u2netp", "u2net_human_seg"]
 
     ap = argparse.ArgumentParser()
 
@@ -101,8 +102,8 @@ def main():
 
     if args.path:
         full_paths = [os.path.abspath(path) for path in args.path]
-        
-        input_paths = [full_paths[0]] 
+
+        input_paths = [full_paths[0]]
         output_path = full_paths[1]
 
         if not os.path.exists(output_path):

+ 2 - 0
src/rembg/cmd/server.py

@@ -45,6 +45,8 @@ def index():
         os.path.expanduser(os.path.join("~", ".u2net")),
     )
     model_choices = [os.path.splitext(os.path.basename(x))[0] for x in set(glob.glob(model_path + "/*"))]
+    if len(model_choices) == 0:
+        model_choices = ["u2net", "u2netp", "u2net_human_seg"]
 
     if model not in model_choices:
         return {"error": f"invalid query param 'model'. Available options are {model_choices}"}, 400