Update app.py

app.py

@@ -1,22 +1,34 @@
 import gradio as gr
 from rembg import remove
 from PIL import Image, ImageDraw
+import numpy as np
 import io
 
-def remove_bg(image, mask=None):
+def remove_bg(image):
     if image is None:
         return None
     image = Image.open(image)
-    
-    # Apply background removal
     output = remove(image)
+    return output
+
+def refine_bg(image, annotations, threshold):
+    if image is None:
+        return None
+    image = Image.open(image).convert("RGBA")
+    mask = Image.new("L", image.size, 0)
+    draw = ImageDraw.Draw(mask)
     
-    # Apply manual mask refinements if provided
-    if mask:
-        mask_image = Image.open(mask).convert("L")
-        output.paste(image, mask=mask_image)
+    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)
     
-    return output
+    refined = Image.composite(image, Image.new("RGBA", image.size, (0, 0, 0, 0)), mask)
+    return refined
 
 # Creating Gradio interface
 with gr.Blocks() as iface:
@@ -24,11 +36,15 @@ with gr.Blocks() as iface:
     
     with gr.Row():
         input_img = gr.Image(type="filepath", label="Upload Image")
-        mask_img = gr.Image(type="filepath", label="Upload Mask (Optional)")
+        output_img = gr.Image(type="pil", label="Output Image")
     
-    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)
     
-    process_btn = gr.Button("Remove Background")
-    process_btn.click(remove_bg, inputs=[input_img, mask_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)
 
 iface.launch()