sidbhasin commited on
Commit
a45dc04
·
verified ·
1 Parent(s): 380a04f

Create app.py

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