Update app.py

app.py

@@ -11,7 +11,7 @@ def remove_bg(image):
     output = remove(image)
     return output
 
-def refine_bg(image, points, threshold):
+def refine_bg(image, points, threshold, mode):
     if image is None:
         return None
     
@@ -20,7 +20,7 @@ def refine_bg(image, points, threshold):
     draw = ImageDraw.Draw(mask)
     
     for point in points:
-        x, y, mode = point
+        x, y = point
         color = 255 if mode == "keep" else 0
         draw.ellipse((x - threshold * 10, y - threshold * 10, x + threshold * 10, y + threshold * 10), fill=color)
     
@@ -28,7 +28,7 @@ def refine_bg(image, points, threshold):
     return refined
 
 with gr.Blocks() as iface:
-    gr.Markdown("# AI Background Remover with Refinement")
+    gr.Markdown("# AI Background Remover")
     
     with gr.Row():
         input_img = gr.Image(type="filepath", label="Upload Image")
@@ -37,20 +37,32 @@ with gr.Blocks() as iface:
     remove_btn = gr.Button("Remove Background")
     remove_btn.click(remove_bg, inputs=input_img, outputs=output_img)
     
-    with gr.Row():
+    with gr.Column():
         refine_btn = gr.Button("Refine")
-        refine_panel = gr.Column(visible=False)
+        refine_options = gr.Row(visible=False)
     
-    with refine_panel:
+    with refine_options:
         gr.Markdown("### Refine Background")
-        refine_editor = gr.AnnotatedImage(label="Tap to Add Points (+ to Keep, - to Remove)")
+        refine_editor = gr.AnnotatedImage(label="Tap to Add Points")
+        keep_btn = gr.Button("+")
+        remove_btn = gr.Button("-")
         threshold_slider = gr.Slider(0.00, 1.00, value=0.5, step=0.01, label="Threshold")
         apply_refine_btn = gr.Button("Apply Refinements")
-        apply_refine_btn.click(refine_bg, inputs=[input_img, refine_editor, threshold_slider], outputs=output_img)
+    
+    refine_mode = gr.State("keep")
     
     def show_refine_section(image):
         return gr.update(visible=True)
     
-    refine_btn.click(show_refine_section, inputs=input_img, outputs=refine_panel)
+    def set_mode_keep():
+        return "keep"
+    
+    def set_mode_remove():
+        return "remove"
+    
+    refine_btn.click(show_refine_section, inputs=input_img, outputs=refine_options)
+    keep_btn.click(set_mode_keep, outputs=refine_mode)
+    remove_btn.click(set_mode_remove, outputs=refine_mode)
+    apply_refine_btn.click(refine_bg, inputs=[input_img, refine_editor, threshold_slider, refine_mode], outputs=output_img)
 
-iface.launch()
+iface.launch()