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

Add more arguments to server

Browse Source
This commit is contained in:
leko 2020-12-29 04:15:24 +08:00 committed by GitHub
parent 58fd707872
commit 3532ed6cf8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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: