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)