File size: 1,411 Bytes
60f6f07 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
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()
|