refactoring

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

View File

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

View File

@ -13,11 +13,8 @@ app = Flask(__name__)
@app.route('/', methods=['GET', 'POST'])
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 = ''
if request.method == 'POST':
if 'file' not in request.files:
return {"error": "missing post form param 'file'"}, 400
@ -34,6 +31,10 @@ def index():
if file_content == '':
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:
return send_file(
BytesIO(remove(file_content, model)),