Update app.py

app.py

@@ -2,68 +2,69 @@ import gradio as gr
 from rembg import remove
 from PIL import Image, ImageDraw
 import numpy as np
+import io
 
-# Store user-selected points
-selected_points = []
-
+# Function to remove background automatically
 def remove_bg(image):
     if image is None:
-        return None, None
-    
-    image = Image.open(image).convert("RGBA")
+        return None
+    image = Image.open(image)
     output = remove(image)
-    return output, output  # Returning twice to keep the output section visible
+    return output
 
-def add_point(image, x, y, mode, threshold):
-    """ Adds a refinement point and applies the changes """
+# Function to refine background using + and - points
+def refine_bg(image, points, threshold, mode):
     if image is None:
         return None
     
-    image = image.convert("RGBA")
+    image = Image.open(image).convert("RGBA")
     mask = Image.new("L", image.size, 0)
     draw = ImageDraw.Draw(mask)
     
-    color = 255 if mode == "keep" else 0
-    radius = int(threshold * 10)
-    draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=color)
+    for point in points:
+        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)
     
     refined = Image.composite(image, Image.new("RGBA", image.size, (0, 0, 0, 0)), mask)
     return refined
 
+# Gradio UI
 with gr.Blocks() as iface:
     gr.Markdown("# AI Background Remover")
-
+    
     with gr.Row():
         input_img = gr.Image(type="filepath", label="Upload Image")
-        output_img = gr.Image(type="pil", label="Output Image", interactive=True)
-
+        output_img = gr.Image(type="pil", label="Output Image")
+    
     remove_btn = gr.Button("Remove Background")
+    remove_btn.click(remove_bg, inputs=input_img, outputs=output_img)
     
     refine_btn = gr.Button("Refine")
-    threshold_slider = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Threshold")
+    refine_options = gr.Column(visible=False)
     
-    with gr.Row():
+    with refine_options:
+        gr.Markdown("### Refine Background")
+        refine_editor = gr.Image(type="filepath", label="Tap to Add Points", interactive=True)
         keep_btn = gr.Button("+ (Keep)")
         remove_btn = gr.Button("- (Remove)")
-
-    apply_refine_btn = gr.Button("Apply Refinements")
+        threshold_slider = gr.Slider(0.00, 1.00, value=0.5, step=0.01, label="Threshold")
+        apply_refine_btn = gr.Button("Apply Refinements")
+    
+    refine_mode = gr.State("keep")
+    
+    def show_refine_section(image):
+        return gr.update(visible=True), image  # Show section and keep image visible
     
-    refine_mode = gr.State("keep")  # Default mode: Keep
-
     def set_mode_keep():
-        return "keep"
+        return "keep", "+ (Keep) Selected"
     
     def set_mode_remove():
-        return "remove"
-    
-    remove_btn.click(remove_bg, inputs=input_img, outputs=[output_img, output_img])
-    keep_btn.click(set_mode_keep, outputs=refine_mode)
-    remove_btn.click(set_mode_remove, outputs=refine_mode)
-
-    def refine_with_clicks(image, evt: gr.SelectData, mode, threshold):
-        """ Captures user clicks and applies AI-assisted refinement. """
-        return add_point(image, evt.index[0], evt.index[1], mode, threshold)
+        return "remove", "- (Remove) Selected"
     
-    output_img.select(refine_with_clicks, inputs=[output_img, refine_mode, threshold_slider], outputs=output_img)
+    refine_btn.click(show_refine_section, inputs=input_img, outputs=[refine_options, refine_editor])
+    keep_btn.click(set_mode_keep, outputs=[refine_mode, keep_btn])
+    remove_btn.click(set_mode_remove, outputs=[refine_mode, remove_btn])
+    apply_refine_btn.click(refine_bg, inputs=[input_img, refine_editor, threshold_slider, refine_mode], outputs=output_img)
 
 iface.launch()