AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-14 11:35:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactoring

Browse Source
This commit is contained in:
Daniel Gatis 2020-09-12 09:29:22 -03:00
parent 26fc23e579
commit dc72f37e5f
2 changed files with 6 additions and 5 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.9", version="1.0.10",
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",

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

@ -13,11 +13,8 @@ app = Flask(__name__)
@app.route('/', methods=['GET', 'POST']) @app.route('/', methods=['GET', 'POST'])
def index(): def index():
model = request.args.get("model", type=str, default="u2net")
if model not in ("u2net", "u2netp"):
return {"error": "invalid query param 'model'"}, 400
file_content = '' file_content = ''
if request.method == 'POST': if request.method == 'POST':
if 'file' not in request.files: if 'file' not in request.files:
return {"error": "missing post form param 'file'"}, 400 return {"error": "missing post form param 'file'"}, 400
@ -34,6 +31,10 @@ def index():
if file_content == '': if file_content == '':
return {"error": "File content is empty"}, 400 return {"error": "File content is empty"}, 400
model = request.args.get("model", type=str, default="u2net")
if model not in ("u2net", "u2netp"):
return {"error": "invalid query param 'model'"}, 400
try: try:
return send_file( return send_file(
BytesIO(remove(file_content, model)), BytesIO(remove(file_content, model)),