File size: 3,141 Bytes
a45dc04 |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import gradio as gr
from transformers import pipeline
import torch
import numpy as np
from PIL import Image
import io
import base64
def remove_background(input_image):
try:
# Initialize the pipeline
segmentor = pipeline("image-segmentation",
model="briaai/RMBG-1.4",
device=-1) # CPU inference
# Process the image
result = segmentor(input_image)
# Return both original and processed images
return result['output_image']
except Exception as e:
raise gr.Error(f"Error processing image: {str(e)}")
# Create the Gradio interface
css = """
.gradio-container {
font-family: 'Segoe UI', sans-serif;
background: linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%);
}
.gr-button {
background: linear-gradient(45deg, #FFD700, #FFA500);
border: none;
color: black;
}
.gr-button:hover {
background: linear-gradient(45deg, #FFA500, #FFD700);
transform: translateY(-2px);
box-shadow: 0 4px 12px rgba(255, 215, 0, 0.3);
}
.gr-form {
background: rgba(255, 255, 255, 0.1);
border-radius: 16px;
padding: 20px;
box-shadow: 0 4px 20px rgba(0, 0, 0, 0.2);
}
.gr-image {
border-radius: 12px;
border: 2px solid rgba(255, 215, 0, 0.3);
}
"""
with gr.Blocks(css=css) as demo:
gr.HTML(
"""
<div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
<h1 style="font-size: 2.5rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
Background Removal Tool
</h1>
<p style="color: #cccccc; font-size: 1.1rem; margin-bottom: 2rem;">
Powered by RMBG V1.4 model from BRIA AI
</p>
</div>
"""
)
with gr.Row():
with gr.Column():
input_image = gr.Image(
label="Upload Image",
type="pil",
tool="upload",
)
with gr.Row():
clear_btn = gr.Button("Clear")
submit_btn = gr.Button("Remove Background", variant="primary")
with gr.Column():
output_image = gr.Image(
label="Result",
type="pil",
)
download_btn = gr.Button("Download Result")
# Event handlers
submit_btn.click(
fn=remove_background,
inputs=[input_image],
outputs=[output_image],
)
clear_btn.click(
lambda: (None, None),
outputs=[input_image, output_image],
)
# Example images
gr.Examples(
examples=[
["example1.jpg"],
["example2.jpg"],
["example3.jpg"],
],
inputs=input_image,
outputs=output_image,
fn=remove_background,
cache_examples=True,
)
# Download functionality
download_btn.click(
lambda x: x,
inputs=[output_image],
outputs=[output_image],
)
demo.launch()
|