Update app.py

app.py

@@ -1,36 +1,107 @@
-import gradio as gr
 import cv2
 import numpy as np
 from PIL import Image
+import gradio as gr
+import os
 
-def refine_edges(image: Image.Image):
-    # Convert to OpenCV format (RGBA)
-    image_np = np.array(image)
-    if image_np.shape[2] == 4:
-        alpha_channel = image_np[:, :, 3]
-    else:
-        return image  # Return original if no alpha channel
-    
-    # Step 1: Blur the alpha channel
-    blurred = cv2.GaussianBlur(alpha_channel, (5, 5), 0)
+# Create examples directory if it doesn't exist
+os.makedirs("examples", exist_ok=True)
 
-    # Step 2: Morphological refinement
-    kernel = np.ones((3,3), np.uint8)
-    refined = cv2.morphologyEx(blurred, cv2.MORPH_OPEN, kernel)
+def refine_edges(image, edge_smoothness=3, blur_radius=2, feather_amount=1):
+    """
+    Refines edges of a transparent PNG image with configurable parameters.
+    
+    Args:
+        image: Input image (PIL Image)
+        edge_smoothness: Intensity of edge refinement (1-5)
+        blur_radius: Kernel size for edge smoothing (1-5)
+        feather_amount: Feathering amount at transparency edges (0-5)
+    
+    Returns:
+        Refined PIL Image
+    """
+    img = image.convert("RGBA")
+    np_img = np.array(img)
+    alpha = np_img[:, :, 3]
+    
+    # Scale parameters
+    blur_kernel = blur_radius * 2 + 1  # 3,5,7,9,11
+    smooth_iterations = edge_smoothness
+    feather_size = feather_amount
+    
+    # Edge smoothing with morphological operations
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+    for _ in range(smooth_iterations):
+        alpha = cv2.morphologyEx(alpha, cv2.MORPH_OPEN, kernel)
+        alpha = cv2.morphologyEx(alpha, cv2.MORPH_CLOSE, kernel)
+    
+    # Gaussian blur for smoother edges
+    alpha = cv2.GaussianBlur(alpha, (blur_kernel, blur_kernel), 0)
+    
+    # Feather edges if enabled
+    if feather_amount > 0:
+        _, mask = cv2.threshold(alpha, 10, 255, cv2.THRESH_BINARY)
+        edges = cv2.Canny(mask, 100, 200)
+        edges = cv2.dilate(edges, np.ones((feather_size, feather_size), np.uint8), iterations=1)
+        alpha_blurred = cv2.GaussianBlur(alpha, (blur_kernel, blur_kernel), 0)
+        alpha = np.where(edges > 0, alpha_blurred, alpha)
+    
+    # Re-normalize alpha
+    alpha = np.clip(alpha, 0, 255).astype(np.uint8)
+    np_img[:, :, 3] = alpha
+    
+    return Image.fromarray(np_img)
 
-    # Step 3: Recombine channels
-    b, g, r = image_np[:, :, 0], image_np[:, :, 1], image_np[:, :, 2]
-    rgba = cv2.merge((b, g, r, refined))
+# Download example images if they don't exist
+example_images = {
+    "hair.png": "https://huggingface.co/spaces/sd-org/remove-bg/resolve/main/examples/hair.png",
+    "furry_animal.png": "https://huggingface.co/spaces/sd-org/remove-bg/resolve/main/examples/animal.png",
+    "glasses.png": "https://huggingface.co/spaces/sd-org/remove-bg/resolve/main/examples/glasses.png"
+}
 
-    return Image.fromarray(rgba)
+for filename, url in example_images.items():
+    if not os.path.exists(f"examples/{filename}"):
+        try:
+            from urllib.request import urlretrieve
+            urlretrieve(url, f"examples/{filename}")
+        except:
+            print(f"Couldn't download example image: {filename}")
 
-iface = gr.Interface(
-    fn=refine_edges,
-    inputs=gr.Image(type="pil", image_mode="RGBA"),
-    outputs="image",
-    title="Edge Refiner",
-    description="Upload a background-removed PNG and get smoother edges!"
-)
+# Create Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    # ✨ Edge Refiner - Clean Up Your Background-Removed Images!
+    Refine the edges of images that have already been background-removed (PNGs with transparency).
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Input Image (Transparent PNG)")
+            edge_smoothness = gr.Slider(1, 5, value=3, step=1, label="Edge Smoothness")
+            blur_radius = gr.Slider(1, 5, value=2, step=1, label="Blur Radius")
+            feather_amount = gr.Slider(0, 5, value=1, step=1, label="Feather Amount")
+            submit_btn = gr.Button("Refine Edges")
+        
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Refined Image")
+    
+    gr.Examples(
+        examples=[
+            ["examples/hair.png", 4, 2, 1],
+            ["examples/furry_animal.png", 3, 3, 2],
+            ["examples/glasses.png", 2, 1, 0]
+        ],
+        inputs=[input_image, edge_smoothness, blur_radius, feather_amount],
+        outputs=output_image,
+        fn=refine_edges,
+        cache_examples=True
+    )
+    
+    submit_btn.click(
+        fn=refine_edges,
+        inputs=[input_image, edge_smoothness, blur_radius, feather_amount],
+        outputs=output_image
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()