File size: 4,616 Bytes
a45dc04
f41e3e9
331d778
 
 
 
a45dc04
 
 
855a559
e67f455
 
 
 
 
 
a45dc04
e67f455
 
 
 
 
 
 
 
 
 
 
 
 
a45dc04
 
 
e67f455
 
 
 
 
 
 
 
 
 
 
 
a45dc04
 
e67f455
 
 
855a559
36588be
e67f455
 
 
 
 
 
 
 
f41e3e9
a45dc04
e67f455
 
 
 
 
 
 
 
 
a45dc04
e67f455
 
 
 
855a559
a45dc04
 
36588be
e67f455
a45dc04
 
e67f455
 
 
a45dc04
e67f455
 
a45dc04
 
 
 
 
e67f455
a45dc04
 
e67f455
a45dc04
e67f455
a45dc04
 
e67f455
 
 
 
 
a45dc04
 
 
36588be
e67f455
a45dc04
e67f455
 
 
 
 
 
 
a45dc04
e67f455
36588be
855a559
e67f455
855a559
36588be
e67f455
 
 
 
 
855a559
e67f455
855a559
f41e3e9
36588be
a45dc04
e67f455
a45dc04
 
 
e67f455
 
a45dc04
 
e67f455
 
 
 
 
 
 
 
 
 
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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import gradio as gr
from transformers import pipeline
import torch
import numpy as np
from PIL import Image
import io

def remove_background(input_image):
    try:
        # Initialize the pipeline with trust_remote_code=True
        segmentor = pipeline(
            "image-segmentation", 
            model="briaai/RMBG-1.4",
            trust_remote_code=True,
            device="cpu"
        )
        
        # Process the image and get result
        result = segmentor(input_image, return_mask=True)
        
        # Convert result to RGBA
        if isinstance(result, Image.Image):
            # Create transparent background
            output = Image.new('RGBA', result.size, (0, 0, 0, 0))
            output.paste(input_image, mask=result)
        else:
            output = result['output_image']
        
        return output

    except Exception as e:
        raise gr.Error(f"Error processing image: {str(e)}")

# Custom theme and styling
theme = gr.themes.Soft(
    primary_hue="gold",
    secondary_hue="orange",
).set(
    body_background_fill="linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%)",
    body_text_color="#ffffff",
    button_primary_background_fill="linear-gradient(45deg, #FFD700, #FFA500)",
    button_primary_text_color="#000000",
    border_color_primary="#FFD700"
)

css = """
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}
.image-container {
    border-radius: 15px !important;
    border: 2px solid rgba(255, 215, 0, 0.3) !important;
    padding: 10px !important;
    background: rgba(255, 255, 255, 0.1) !important;
    transition: transform 0.3s ease !important;
}
.image-container:hover {
    transform: translateY(-5px) !important;
}
.gr-button {
    min-width: 200px !important;
    height: 45px !important;
    font-size: 16px !important;
    margin: 10px !important;
    transition: all 0.3s ease !important;
}
.gr-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 5px 15px rgba(255, 215, 0, 0.3) !important;
}
.footer {
    text-align: center;
    margin-top: 20px;
    color: #666;
}
"""

# Create Gradio interface
with gr.Blocks(theme=theme, css=css) as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin-bottom: 2rem;">
            <h1 style="font-size: 3rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
                AI Background Remover Pro
            </h1>
            <p style="color: #cccccc; font-size: 1.2rem;">
                Remove backgrounds instantly using advanced AI technology
            </p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Your Image",
                type="pil",
                elem_classes="image-container"
            )
            
            with gr.Row():
                clear_btn = gr.Button("Clear", variant="secondary")
                process_btn = gr.Button("Remove Background", variant="primary")
                download_btn = gr.Button("Download", variant="primary", visible=False)

        with gr.Column():
            output_image = gr.Image(
                label="Result",
                type="pil",
                elem_classes="image-container"
            )
            
    # Status message
    status_msg = gr.Textbox(
        label="Status",
        placeholder="Ready to process your image...",
        interactive=False
    )

    # Event handlers
    def process_and_update(image):
        if image is None:
            return None, "Please upload an image first", gr.Button.update(visible=False)
        try:
            result = remove_background(image)
            return (
                result, 
                "✨ Background removed successfully!", 
                gr.Button.update(visible=True)
            )
        except Exception as e:
            return None, f"❌ Error: {str(e)}", gr.Button.update(visible=False)

    process_btn.click(
        fn=process_and_update,
        inputs=[input_image],
        outputs=[output_image, status_msg, download_btn],
    )
    
    clear_btn.click(
        fn=lambda: (None, None, "Ready to process your image...", gr.Button.update(visible=False)),
        outputs=[input_image, output_image, status_msg, download_btn],
    )

    gr.HTML(
        """
        <div class="footer">
            <p>Powered by BRIA AI's RMBG V1.4 Model</p>
        </div>
        """
    )

# Launch the app
demo.launch()