import gradio as gr
import cv2
import numpy as np
from PIL import Image
import io, base64, json
from fastapi import FastAPI, Request
from fastapi.responses import Response
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

@app.post("/api/erase")
async def api_erase(request: Request):
    body = await request.json()
    img = Image.open(io.BytesIO(base64.b64decode(body['image'])))
    msk = Image.open(io.BytesIO(base64.b64decode(body['mask'])))
    img_np = np.array(img.convert('RGB'))
    msk_np = np.array(msk.convert('L'))
    _, msk_np = cv2.threshold(msk_np, 127, 255, cv2.THRESH_BINARY)
    result = cv2.inpaint(img_np, msk_np, 3, cv2.INPAINT_TELEA)
    buf = io.BytesIO()
    Image.fromarray(result).save(buf, 'PNG')
    return Response(content=buf.getvalue(), media_type='image/png')

def inpaint_ui(image, mask):
    img = np.array(image.convert('RGB'))
    msk = np.array(mask.convert('L'))
    _, msk = cv2.threshold(msk, 127, 255, cv2.THRESH_BINARY)
    result = cv2.inpaint(img, msk, 3, cv2.INPAINT_TELEA)
    return Image.fromarray(result)

css = 'footer{display:none}'
with gr.Blocks(title='Magic Eraser', css=css) as demo:
    gr.Markdown('# ✨ Magic Eraser')
    with gr.Row():
        with gr.Column():
            image = gr.Image(label='Photo', type='pil')
            mask = gr.Image(label='Mask (white = area to remove)', type='pil')
            btn = gr.Button('✨ Erase Object')
        with gr.Column():
            result = gr.Image(label='Result', type='pil')
    btn.click(fn=inpaint_ui, inputs=[image, mask], outputs=[result])
    gr.Markdown('Draw the area to remove in **white** on the mask image.')

gr.mount_gradio_app(app, demo, path="/")