from flask import Flask, request, jsonify
from gradio_client import Client, handle_file
from PIL import Image
import requests
import io
import base64
import numpy as np
from io import BytesIO

app = Flask(__name__)

client = Client("ECCV2022/dis-background-removal")

def image_to_data_url(img: Image.Image) -> str:
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{img_str}"
def base64_to_image(data_url):
    header, encoded = data_url.split(",", 1)
    binary_data = base64.b64decode(encoded)
    image = Image.open(BytesIO(binary_data))
    return image

import tempfile

@app.route("/remove_background", methods=["POST"])
def remove_background():
    data = request.json
    data_url = data.get("image_url")

    if not data_url:
        return jsonify({"error": "image_url is required"}), 400

    try:
        # base64データURLから画像を取得して一時ファイルに保存
        input_image = base64_to_image(data_url).convert("RGB")
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            input_image.save(tmp, format="PNG")
            temp_path = tmp.name

        # Gradioのモデルに画像を渡して背景除去を実行
        result = client.predict(
            image=handle_file(temp_path),
            api_name="/predict"
        )

        # 結果として返される2つの画像パスを取得
        image_path1, image_path2 = result
        depth_image = Image.open(image_path2)
        processed_image = Image.open(image_path1)

        return jsonify({
            "depth_image": image_to_data_url(depth_image),
            "semi_transparent_image": image_to_data_url(processed_image)
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=7860)