| 15 | |
| 16 | @app.route("/", methods=["GET", "POST"]) |
| 17 | def index(): |
| 18 | file_content = "" |
| 19 | |
| 20 | if request.method == "POST": |
| 21 | if "file" not in request.files: |
| 22 | return {"error": "missing post form param 'file'"}, 400 |
| 23 | |
| 24 | file_content = request.files["file"].read() |
| 25 | |
| 26 | if request.method == "GET": |
| 27 | url = request.args.get("url", type=str) |
| 28 | if url is None: |
| 29 | return {"error": "missing query param 'url'"}, 400 |
| 30 | |
| 31 | file_content = urlopen(unquote_plus(url)).read() |
| 32 | |
| 33 | if file_content == "": |
| 34 | return {"error": "File content is empty"}, 400 |
| 35 | |
| 36 | alpha_matting = "a" in request.values |
| 37 | af = request.values.get("af", type=int, default=240) |
| 38 | ab = request.values.get("ab", type=int, default=10) |
| 39 | ae = request.values.get("ae", type=int, default=10) |
| 40 | az = request.values.get("az", type=int, default=1000) |
| 41 | mt = request.values.get("mt", type=int, default=None) |
| 42 | |
| 43 | model = request.args.get("model", type=str, default="u2net") |
| 44 | model_path = os.environ.get( |
| 45 | "U2NETP_PATH", |
| 46 | os.path.expanduser(os.path.join("~", ".u2net")), |
| 47 | ) |
| 48 | model_choices = [os.path.splitext(os.path.basename(x))[0] for x in set(glob.glob(model_path + "/*"))] |
| 49 | if len(model_choices) == 0: |
| 50 | model_choices = ["u2net", "u2netp", "u2net_human_seg"] |
| 51 | |
| 52 | if model not in model_choices: |
| 53 | return {"error": f"invalid query param 'model'. Available options are {model_choices}"}, 400 |
| 54 | |
| 55 | try: |
| 56 | return send_file( |
| 57 | BytesIO( |
| 58 | remove( |
| 59 | file_content, |
| 60 | model_name=model, |
| 61 | alpha_matting=alpha_matting, |
| 62 | alpha_matting_foreground_threshold=af, |
| 63 | alpha_matting_background_threshold=ab, |
| 64 | alpha_matting_erode_structure_size=ae, |
| 65 | alpha_matting_base_size=az, |
| 66 | mask_threshold=mt, |
| 67 | ) |
| 68 | ), |
| 69 | mimetype="image/png", |
| 70 | ) |
| 71 | except Exception as e: |
| 72 | app.logger.exception(e, exc_info=True) |
| 73 | return {"error": "oops, something went wrong!"}, 500 |
| 74 | |