AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-18 03:55:54 +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( setup(
name="rembg", name="rembg",
version="1.0.25", version="1.0.26",
description="Remove image background", description="Remove image background",
long_description=long_description, long_description=long_description,
long_description_content_type="text/markdown", 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")), 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 + "/*"))] 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() ap = argparse.ArgumentParser()

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

@ -45,6 +45,8 @@ def index():
os.path.expanduser(os.path.join("~", ".u2net")), 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 + "/*"))] 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: if model not in model_choices:
return {"error": f"invalid query param 'model'. Available options are {model_choices}"}, 400 return {"error": f"invalid query param 'model'. Available options are {model_choices}"}, 400