Update app.py

app.py

@@ -11,28 +11,17 @@ def remove_bg(image):
     output = remove(image)
     return output
 
-def refine_bg(image, annotations, threshold):
-    if image is None:
+def refine_bg(image, mask, threshold):
+    if image is None or mask is None:
         return None
     image = Image.open(image).convert("RGBA")
-    mask = Image.new("L", image.size, 0)
-    draw = ImageDraw.Draw(mask)
-    
-    for annotation in annotations:
-        points = annotation["points"]
-        label = annotation["label"]
-        
-        if label == "keep":
-            draw.polygon(points, fill=255)
-        else:  # label == "remove"
-            draw.polygon(points, fill=0)
-    
+    mask = Image.open(mask).convert("L")
     refined = Image.composite(image, Image.new("RGBA", image.size, (0, 0, 0, 0)), mask)
     return refined
 
 # Creating Gradio interface
 with gr.Blocks() as iface:
-    gr.Markdown("# AI Background Remover with Refinement")
+    gr.Markdown("# AI Background Remover with Inline Refinement")
     
     with gr.Row():
         input_img = gr.Image(type="filepath", label="Upload Image")
@@ -41,10 +30,10 @@ with gr.Blocks() as iface:
     remove_btn = gr.Button("Remove Background")
     remove_btn.click(remove_bg, inputs=input_img, outputs=output_img)
     
-    with gr.Accordion("Edit Background Refinement", open=False):
-        annotator = gr.AnnotatedImage(label="Refine Background")
-        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, annotator, threshold_slider], outputs=output_img)
+    gr.Markdown("### Refine Background Directly on Image")
+    mask_editor = gr.Image(type="filepath", label="Draw on Image (Use + for Keep, - for Remove)", tool="sketch")
+    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)
 
 iface.launch()