Added default parameters to the function signature, Updated examples with proper values for all 7 parameters
Browse files
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
|
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 |
-
#
|
|
|
|
|
|
|
|
|
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, **
|
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 |
-
|
129 |
-
True, # transparent_bg
|
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=
|
143 |
-
gr.Checkbox(label="Transparent Background", value=
|
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
|