import gradio as gr
from rembg import remove
from PIL import Image
from flask import Flask, request, jsonify

app = Flask(__name__)

def remove_background(input_image):
    output_image = remove(input_image)
    return output_image

inputs = gr.inputs.Image()
outputs = gr.outputs.Image(type='pil')

interface = gr.Interface(fn=remove_background, inputs=inputs, outputs=outputs, title="Remove Background", description="This App removes the background from an image.", api="flask")
interface.launch()

@app.route("/api/remove-background", methods=["POST"])
def remove_background_api():
    input_image = Image.open(request.files["input"])
    output_image = remove_background(input_image)
    return jsonify({"output": output_image})

# Custom HTML block with API instructions
interface.interface_html = """
    <div style="padding: 1rem; margin: 1rem 0; border: 1px solid black;">
        <h3>API Instructions:</h3>
        <ol>
            <li>Send a POST request to the '/api/remove-background' endpoint with an image file in the 'input' field of the request body.</li>
            <li>The output image will be returned as a byte string in the 'output' field of the JSON response.</li>
        </ol>
    </div>
"""