Create app.py

app.py

@@ -0,0 +1,38 @@
+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()