Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from rembg import remove
-from PIL import Image
+from PIL import Image, ImageDraw
 import numpy as np
 import io
 
@@ -11,11 +11,19 @@ def remove_bg(image):
     output = remove(image)
     return output
 
-def refine_bg(image, mask, threshold):
-    if image is None or mask is None:
+def refine_bg(image, points, threshold):
+    if image is None or points is None:
         return None
     image = Image.open(image).convert("RGBA")
-    mask = Image.open(mask).convert("L")
+    mask = Image.new("L", image.size, 0)
+    draw = ImageDraw.Draw(mask)
+    
+    for point in points:
+        x, y, action = point  # action: 1 for keep (+), 0 for remove (-)
+        radius = int(threshold * 50)  # Scale the threshold to radius size
+        color = 255 if action == 1 else 0  # White for keep, Black for remove
+        draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=color)
+    
     refined = Image.composite(image, Image.new("RGBA", image.size, (0, 0, 0, 0)), mask)
     return refined
 
@@ -30,10 +38,17 @@ with gr.Blocks() as iface:
     remove_btn = gr.Button("Remove Background")
     remove_btn.click(remove_bg, inputs=input_img, outputs=output_img)
     
-    gr.Markdown("### Refine Background Directly on Image")
-    mask_editor = gr.Image(type="filepath", label="Upload Mask (Use Black for Remove, White for Keep)")
-    threshold_slider = gr.Slider(0.00, 1.00, value=0.5, step=0.01, label="Threshold")
-    refine_btn = gr.Button("Apply Refinements")
-    refine_btn.click(refine_bg, inputs=[input_img, mask_editor, threshold_slider], outputs=output_img)
+    with gr.Column(visible=False) as refine_section:
+        gr.Markdown("### Refine Background")
+        refine_editor = gr.Image(type="filepath", label="Tap to Add Points (+ to Keep, - to Remove)", tool="editor")
+        threshold_slider = gr.Slider(0.00, 1.00, value=0.5, step=0.01, label="Threshold")
+        refine_btn = gr.Button("Apply Refinements")
+        refine_btn.click(refine_bg, inputs=[input_img, refine_editor, threshold_slider], outputs=output_img)
+    
+    def show_refine_section(image):
+        return gr.update(visible=True)
+    
+    refine_toggle_btn = gr.Button("Refine")
+    refine_toggle_btn.click(show_refine_section, inputs=input_img, outputs=refine_section)
 
-iface.launch()
+iface.launch()