AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-17 06:46:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix invalid model fresh install

Browse Source
This commit is contained in:
Daniel Gatis 2021-04-04 20:07:11 -03:00
parent ca4be25f90
commit b07192d16f
3 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
setup.py
View File

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

3
src/rembg/cmd/cli.py
View File

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

2
src/rembg/cmd/server.py
View File

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