浏览代码

Merge pull request #5 from rcrvano/master

Support for sending file with POST request
Daniel Gatis 5 年之前
父节点
当前提交
279a50bdf8
共有 2 个文件被更改,包括 26 次插入4 次删除
  1. 8 0
      README.md
  2. 18 4
      src/rembg/cmd/server.py

+ 8 - 0
README.md

@@ -72,6 +72,14 @@ Open your browser to
     http://localhost:5000?url=http://image.png
 ```
 
+Also you can send the file as a FormData (multipart/form-data):
+```
+    <form action='http://localhost:5000' method='post' enctype='multipart/form-data'>
+      <input type=file name=file>
+      <input type=submit value=Upload>
+    </form>
+```
+
 ### Usage as a library
 
 In `app.py`

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

@@ -11,18 +11,32 @@ from ..bg import remove
 app = Flask(__name__)
 
 
[email protected]('/', 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
 
-    url = request.args.get("url", type=str)
-    if url is None:
-        return {"error": "missing query param 'url'"}, 400
+    file_content = ''
+    if request.method == 'POST':
+        if 'file' not in request.files:
+            return {"error": "missing post form param 'file'"}, 400
+
+        file_content = request.files['file'].read();
+
+    if request.method == 'GET':
+        url = request.args.get("url", type=str)
+        if url is None:
+            return {"error": "missing query param 'url'"}, 400
+
+        file_content = urlopen(unquote_plus(url)).read();
+
+    if file_content == '':
+        return {"error": "File content is empty"}, 400
 
     try:
         return send_file(
-            BytesIO(remove(urlopen(unquote_plus(url)).read(), model)),
+            BytesIO(remove(file_content, model)),
             mimetype="image/png",
         )
     except Exception as e: