import gradio as gr
from transformers import pipeline
from PIL import Image

# Initialize the pipeline for image segmentation
pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)


def process_image(image_path, include_background, background_color):
    # Apply the mask on the input image
    pillow_image = pipe(image_path)

    if include_background:
        # Create a background of the specified color
        background = Image.new("RGB", pillow_image.size, background_color)
        # Combine the background and the image
        background.paste(pillow_image, (0, 0), pillow_image)
        return background
    else:
        return pillow_image


# Define Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Background Removal with Optional Custom Background")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Upload Image")
            include_background = gr.Checkbox(label="Include Background", value=False)
            background_color = gr.ColorPicker(label="Background Color", value="#FFFFFF")
        with gr.Column():
            output_image = gr.Image(label="Processed Image")

            submit_button = gr.Button("Process Image")
    submit_button.click(process_image, inputs=[image_input, include_background, background_color], outputs=output_image)

# Launch the Gradio app
demo.launch()