Bläddra i källkod

Merge pull request #56 from lekoOwO/patch-1

Add more arguments to server
Daniel Gatis 4 år sedan
förälder
incheckning
8509f0448e
1 ändrade filer med 15 tillägg och 1 borttagningar
  1. 15 1
      src/rembg/cmd/server.py

+ 15 - 1
src/rembg/cmd/server.py

@@ -31,13 +31,27 @@ def index():
     if file_content == "":
     if file_content == "":
         return {"error": "File content is empty"}, 400
         return {"error": "File content is empty"}, 400
 
 
+    alpha_matting = "a" in request.values
+    af = request.values.get("af", type=int, default=240)
+    ab = request.values.get("ab", type=int, default=10)
+    ae = request.values.get("ae", type=int, default=10)
+    
     model = request.args.get("model", type=str, default="u2net")
     model = request.args.get("model", type=str, default="u2net")
     if model not in ("u2net", "u2netp"):
     if model not in ("u2net", "u2netp"):
         return {"error": "invalid query param 'model'"}, 400
         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_name=model,
+                    alpha_matting=alpha_matting,
+                    alpha_matting_foreground_threshold=af,
+                    alpha_matting_background_threshold=ab,
+                    alpha_matting_erode_structure_size=ae
+                )
+            ),
             mimetype="image/png",
             mimetype="image/png",
         )
         )
     except Exception as e:
     except Exception as e: