sidbhasin commited on
Commit
f41e3e9
1 Parent(s): 37e3f85

Update app.py

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