Update app.py

app.py

@@ -3,39 +3,32 @@ from rembg import remove
 from PIL import Image, ImageDraw
 import io
 
-def remove_bg(image, points):
+def remove_bg(image, mask=None):
     if image is None:
         return None
+    image = Image.open(image)
     
-    # Load image
-    image = Image.open(image).convert("RGBA")
-    
-    # Convert points format
-    keep_points = [tuple(p["coordinates"]) for p in points if p["label"] == "keep"]
-    remove_points = [tuple(p["coordinates"]) for p in points if p["label"] == "remove"]
-    
-    # Apply automatic background removal
+    # Apply background removal
     output = remove(image)
     
-    # Apply manual refinement (simple demonstration using circles)
-    draw = ImageDraw.Draw(output)
-    for point in keep_points:
-        draw.ellipse([point[0]-5, point[1]-5, point[0]+5, point[1]+5], fill=(255, 255, 255, 255))
-    for point in remove_points:
-        draw.ellipse([point[0]-5, point[1]-5, point[0]+5, point[1]+5], fill=(0, 0, 0, 0))
+    # Apply manual mask refinements if provided
+    if mask:
+        mask_image = Image.open(mask).convert("L")
+        output.paste(image, mask=mask_image)
     
     return output
 
 # Creating Gradio interface
-iface = gr.Interface(
-    fn=remove_bg,
-    inputs=[
-        gr.Image(type="filepath"),
-        gr.AnnotatedImage(choices=["keep", "remove"], label="Refine Background")
-    ],
-    outputs=gr.Image(type="pil"),
-    title="AI Background Remover with Refinement",
-    description="Upload an image, remove its background, and refine the result using + (keep) and - (remove) points!",
-)
+with gr.Blocks() as iface:
+    gr.Markdown("# AI Background Remover with Refinement")
+    
+    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")
+    
+    process_btn = gr.Button("Remove Background")
+    process_btn.click(remove_bg, inputs=[input_img, mask_img], outputs=output_img)
 
 iface.launch()