import gradio as gr
from rembg import remove, new_session
from PIL import Image

session = new_session()

def remove_bg(input_img, crop=False):
    img = remove(input_img, session=session)
    if crop:
        bbox = img.getchannel("A").getbbox()
        if bbox:
            img = img.crop(bbox)
    return img

def composite(fg_img, bg_img, crop=False, x=0, y=0, scale=1.0):
    fg = remove(fg_img, session=session)
    if crop:
        bbox = fg.getchannel("A").getbbox()
        if bbox:
            fg = fg.crop(bbox)
    fw, fh = fg.size
    fg = fg.resize((int(fw * scale), int(fh * scale)))
    bg = bg_img.convert("RGBA")
    cx = (bg.width - fg.width) // 2 + x
    cy = (bg.height - fg.height) // 2 + y
    bg.paste(fg, (cx, cy), fg)
    return bg.convert("RGB")

with gr.Blocks(title="Background Remover", fill_width=True) as demo:
    gr.Markdown("# Background Remover")
    with gr.Tabs():
        with gr.Tab("Remove Background"):
            gr.Markdown("Remove the background from an image, leaving a transparent result.")
            with gr.Row():
                with gr.Column():
                    inp = gr.Image(label="Input Image", type="pil")
                    crop1 = gr.Checkbox(label="Crop to content", info="Remove transparent edges so the foreground fits tightly")
                    btn1 = gr.Button("Remove Background", variant="primary")
                with gr.Column():
                    out1 = gr.Image(label="Processed Image")
            btn1.click(fn=remove_bg, inputs=[inp, crop1], outputs=out1)
        with gr.Tab("Composite"):
            gr.Markdown("Remove the background from a foreground object and place it onto a background. Adjust position and scale with the controls, then use the built-in editor to brush/erase for final polish.")
            with gr.Row():
                with gr.Column():
                    fg_inp = gr.Image(label="Foreground", type="pil")
                    bg_inp = gr.Image(label="Background", type="pil")
                    crop2 = gr.Checkbox(label="Crop foreground to content")
                    x_off = gr.Slider(-500, 500, 0, step=5, label="X Offset")
                    y_off = gr.Slider(-500, 500, 0, step=5, label="Y Offset")
                    scale = gr.Slider(0.1, 2.0, 1.0, step=0.05, label="Scale")
                    btn2 = gr.Button("Composite", variant="primary")
                with gr.Column():
                    out2 = gr.ImageEditor(label="Result — use brush/eraser to polish")
            btn2.click(fn=composite, inputs=[fg_inp, bg_inp, crop2, x_off, y_off, scale], outputs=out2)

if __name__ == "__main__":
    demo.launch(mcp_server=True)