rdjarbeng commited on
Commit
cb762b0
·
1 Parent(s): bf78000

Added default parameters to the function signature, Updated examples with proper values for all 7 parameters

Browse files
Files changed (1) hide show
  1. app.py +27 -16
app.py CHANGED
@@ -51,12 +51,14 @@ MODEL_OPTIONS = {
51
  }
52
 
53
  def hex_to_rgba(hex_color):
 
 
54
  hex_color = hex_color.lstrip('#')
55
  if len(hex_color) == 6:
56
  hex_color += 'FF' # Add full opacity if no alpha is provided
57
  return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4, 6))
58
 
59
- def remove_background(input_path, bg_color, transparent_bg, model_choice, alpha_matting, post_process_mask, only_mask):
60
  try:
61
  # Log user interaction (minimal)
62
  log_user_interaction()
@@ -76,30 +78,37 @@ def remove_background(input_path, bg_color, transparent_bg, model_choice, alpha_
76
  session = new_session(model_name) if model_name else None
77
 
78
  # Use transparent background if selected, otherwise use color
79
- bg_color_rgba = None if transparent_bg else (hex_to_rgba(bg_color) if bg_color else None)
80
 
81
  # Prepare additional options
82
  remove_kwargs = {
83
  "session": session,
84
- "bgcolor": bg_color_rgba,
85
- "alpha_matting": alpha_matting,
86
- "post_process_mask": post_process_mask,
87
- "only_mask": only_mask
88
  }
89
-
90
- # Add alpha matting parameters if enabled
 
 
 
 
91
  if alpha_matting:
92
  remove_kwargs.update({
 
93
  "alpha_matting_foreground_threshold": 270,
94
  "alpha_matting_background_threshold": 20,
95
  "alpha_matting_erode_size": 11
96
  })
 
 
 
 
 
 
97
 
98
  # Convert PIL Image to numpy array
99
  input_array = np.array(input_image)
100
 
101
  # Use the remove function
102
- output_array = remove(input_array, **{k: v for k, v in remove_kwargs.items() if v is not None})
103
 
104
  # Convert numpy array back to PIL Image
105
  output_image = Image.fromarray(output_array)
@@ -122,12 +131,13 @@ def remove_background(input_path, bg_color, transparent_bg, model_choice, alpha_
122
  logging.error(f"An error occurred: {e}")
123
  return None
124
 
 
125
  examples = [
126
  [
127
  'scifi_man1.jpg', # input_path
128
- None, # bg_color
129
- True, # transparent_bg (most common use case)
130
- "", # model_choice
131
  False, # alpha_matting
132
  False, # post_process_mask
133
  False # only_mask
@@ -139,12 +149,13 @@ iface = gr.Interface(
139
  fn=remove_background,
140
  inputs=[
141
  gr.Image(type="filepath", label="Input Image"),
142
- gr.ColorPicker(label="Background Color (ignored if transparent is selected)", value=None),
143
- gr.Checkbox(label="Transparent Background", value=False),
144
  gr.Dropdown(
145
  choices=[""] + [f"{k} | {v}" for k, v in MODEL_OPTIONS.items() if k != ""],
146
  label="Model Selection",
147
- value=""
 
148
  ),
149
  gr.Checkbox(label="Enable Alpha Matting", value=False),
150
  gr.Checkbox(label="Post-Process Mask", value=False),
@@ -160,4 +171,4 @@ iface = gr.Interface(
160
  )
161
 
162
  if __name__ == "__main__":
163
- iface.launch()
 
51
  }
52
 
53
  def hex_to_rgba(hex_color):
54
+ if not hex_color:
55
+ return None
56
  hex_color = hex_color.lstrip('#')
57
  if len(hex_color) == 6:
58
  hex_color += 'FF' # Add full opacity if no alpha is provided
59
  return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4, 6))
60
 
61
+ def remove_background(input_path, bg_color=None, transparent_bg=True, model_choice="", alpha_matting=False, post_process_mask=False, only_mask=False):
62
  try:
63
  # Log user interaction (minimal)
64
  log_user_interaction()
 
78
  session = new_session(model_name) if model_name else None
79
 
80
  # Use transparent background if selected, otherwise use color
81
+ bg_color_rgba = None if transparent_bg else hex_to_rgba(bg_color)
82
 
83
  # Prepare additional options
84
  remove_kwargs = {
85
  "session": session,
 
 
 
 
86
  }
87
+
88
+ # Only add bgcolor if we have one
89
+ if bg_color_rgba is not None:
90
+ remove_kwargs["bgcolor"] = bg_color_rgba
91
+
92
+ # Add other parameters
93
  if alpha_matting:
94
  remove_kwargs.update({
95
+ "alpha_matting": True,
96
  "alpha_matting_foreground_threshold": 270,
97
  "alpha_matting_background_threshold": 20,
98
  "alpha_matting_erode_size": 11
99
  })
100
+
101
+ if post_process_mask:
102
+ remove_kwargs["post_process_mask"] = True
103
+
104
+ if only_mask:
105
+ remove_kwargs["only_mask"] = True
106
 
107
  # Convert PIL Image to numpy array
108
  input_array = np.array(input_image)
109
 
110
  # Use the remove function
111
+ output_array = remove(input_array, **remove_kwargs)
112
 
113
  # Convert numpy array back to PIL Image
114
  output_image = Image.fromarray(output_array)
 
131
  logging.error(f"An error occurred: {e}")
132
  return None
133
 
134
+ # Fixed examples with proper values for all inputs
135
  examples = [
136
  [
137
  'scifi_man1.jpg', # input_path
138
+ "#FFFFFF", # bg_color (white)
139
+ True, # transparent_bg
140
+ "", # model_choice (empty string, which is in the choices)
141
  False, # alpha_matting
142
  False, # post_process_mask
143
  False # only_mask
 
149
  fn=remove_background,
150
  inputs=[
151
  gr.Image(type="filepath", label="Input Image"),
152
+ gr.ColorPicker(label="Background Color (ignored if transparent is selected)", value="#FFFFFF"),
153
+ gr.Checkbox(label="Transparent Background", value=True),
154
  gr.Dropdown(
155
  choices=[""] + [f"{k} | {v}" for k, v in MODEL_OPTIONS.items() if k != ""],
156
  label="Model Selection",
157
+ value="", # Changed from empty to match choices
158
+ allow_custom_value=False
159
  ),
160
  gr.Checkbox(label="Enable Alpha Matting", value=False),
161
  gr.Checkbox(label="Post-Process Mask", value=False),
 
171
  )
172
 
173
  if __name__ == "__main__":
174
+ iface.launch(share=True) # Added share=True for Hugging Face deployment