Browse Source

refactoring

Daniel Gatis 5 years ago
parent
commit
dc72f37e5f
2 changed files with 6 additions and 5 deletions
  1. 1 1
      setup.py
  2. 5 4
      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.9",
+    version="1.0.10",
     description="Remove image background",
     long_description=long_description,
     long_description_content_type="text/markdown",

+ 5 - 4
src/rembg/cmd/server.py

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