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

def remove_background(input_image):
    """
    Takes a PIL Image, removes the background using rembg, 
    and returns the processed image with a transparent background.
    """
    if input_image is None:
        return None
    
    # The remove() function from rembg handles the heavy lifting
    output_image = remove(input_image)
    return output_image

# Build the Gradio UI
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # ✂️ AI Background Remover
        Upload any image to instantly remove its background. This tool uses the open-source `rembg` library and runs efficiently on CPU.
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Original Image")
            submit_btn = gr.Button("Remove Background", variant="primary")
            clear_btn = gr.ClearButton(components=[input_img])
            
        with gr.Column():
            # Output format is explicitly set to png to preserve transparency
            output_img = gr.Image(type="pil", label="Transparent Result", format="png")
            
    # Connect the button to the function
    submit_btn.click(
        fn=remove_background,
        inputs=input_img,
        outputs=output_img
    )
    
    # Connect the clear button to the output as well
    clear_btn.add(components=[output_img])

# Launch the app (Hugging Face Spaces automatically handles the port and server)
if __name__ == "__main__":
    demo.launch()