huedaya's picture
up
f6a1d78
from flask import Flask, jsonify, request
import cv2
from rembg.session_factory import new_session
from rembg import remove
import requests
import base64
import numpy as np
import streamlit as st
app = Flask(__name__)
@app.route("/")
def indexApi():
return jsonify({"output": "okay"})
@app.route("/remove", methods = ['POST'])
def removeApi():
image_url = request.form.get("image_url")
key = request.form.get("key")
model = request.form.get("model")
if (model == None):
model = "u2netp"
# # reject if key not the same
# apiKey = st.secrets["Api-Key"]
# if apiKey != key:
# return jsonify({
# "image_url": image_url,
# "model": model,
# "result": "Invalid API Key",
# }), 400
# response = requests.get(image_url)
encoded_data = image_url.split(',')[1]
arr = np.fromstring(base64.b64decode(encoded_data), np.uint8)
image = cv2.imdecode(arr, cv2.IMREAD_COLOR)
result = remove(image, session=new_session(model))
#encode result
base64_string = base64.b64encode(cv2.imencode('.png', result)[1]).decode()
return jsonify({
"model": model,
"result": base64_string,
})
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)