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

def replace_bg(input_img, bg_img):
    # Convert input image to RGBA
    input_img = input_img.convert("RGBA")
    
    # Remove background
    no_bg_image = remove(input_img)
    
    # If no background image provided, return transparent background
    if bg_img is None:
        return no_bg_image
    
    # Resize background to match foreground size
    bg_img = bg_img.convert("RGBA").resize(no_bg_image.size)
    
    # Combine foreground and background
    result_image = Image.alpha_composite(bg_img, no_bg_image)
    return result_image

# Create the Gradio interface
demo = gr.Interface(
    fn=replace_bg,
    inputs=[
        gr.Image(type="pil", label="📷 Upload Your Image"),
        gr.Image(type="pil", label="🎨 Upload Background Image (Optional)")
    ],
    outputs=gr.Image(type="pil", label="✨ Final Result"),
    title="🎨 Professional Background Replacer",
    description="Upload any image and replace its background with a new one. Leave background empty for transparent result."
)

# Launch the application
if __name__ == "__main__":
    demo.launch()