Update app.py

app.py

@@ -3,103 +3,79 @@ import numpy as np
 from PIL import Image
 import gradio as gr
 import os
-from io import BytesIO
-import requests
 
-def refine_edges(image, edge_smoothness=3, blur_radius=2, feather_amount=1):
+def refine_edges(image, edge_smoothness=3, blur_radius=2, feather_amount=1, threshold=0.1):
     """
-    Refines edges of a transparent PNG image with configurable parameters.
+    Enhanced edge refinement that specifically targets leftover background pixels
     """
     img = image.convert("RGBA")
     np_img = np.array(img)
-    alpha = np_img[:, :, 3]
+    r, g, b, a = cv2.split(np_img)
     
-    # Scale parameters
+    # Convert parameters
     blur_kernel = blur_radius * 2 + 1
     smooth_iterations = edge_smoothness
-    feather_size = feather_amount
     
-    # Edge smoothing
+    # 1. Create a strict alpha mask (remove semi-transparent pixels)
+    _, strict_alpha = cv2.threshold(a, 254, 255, cv2.THRESH_BINARY)
+    
+    # 2. Find the "edge zone" (area between strict alpha and original alpha)
+    edge_zone = cv2.bitwise_xor(a, strict_alpha)
+    
+    # 3. Process only the edge zone
     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)
+        edge_zone = cv2.morphologyEx(edge_zone, cv2.MORPH_OPEN, kernel)
+        edge_zone = cv2.morphologyEx(edge_zone, cv2.MORPH_CLOSE, kernel)
     
-    # Gaussian blur
-    alpha = cv2.GaussianBlur(alpha, (blur_kernel, blur_kernel), 0)
+    # 4. Apply blur only to edge zone
+    blurred_edge = cv2.GaussianBlur(edge_zone, (blur_kernel, blur_kernel), 0)
     
-    # Feather edges
+    # 5. Feathering with thresholding to remove leftover bg
     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)
+        edge_mask = (blurred_edge > threshold * 255).astype(np.uint8) * 255
+        edge_mask = cv2.erode(edge_mask, np.ones((feather_amount, feather_amount), np.uint8))
+        final_edge = cv2.GaussianBlur(edge_mask, (blur_kernel, blur_kernel), 0)
+    else:
+        final_edge = blurred_edge
     
-    alpha = np.clip(alpha, 0, 255).astype(np.uint8)
-    np_img[:, :, 3] = alpha
+    # 6. Combine with strict alpha
+    new_alpha = cv2.bitwise_or(strict_alpha, final_edge)
     
-    return Image.fromarray(np_img)
-
-def download_example_images():
-    """Download example images from alternative sources"""
-    example_images = {
-        "hair.png": "https://i.imgur.com/JQJQJQJ.png",  # Replace with actual URL
-        "furry_animal.png": "https://i.imgur.com/ANIMAL.png",
-        "glasses.png": "https://i.imgur.com/GLASSES.png"
-    }
+    # 7. Remove color information from transparent areas
+    r = r * (new_alpha > 0)
+    g = g * (new_alpha > 0)
+    b = b * (new_alpha > 0)
     
-    os.makedirs("examples", exist_ok=True)
+    # Recombine channels
+    result = cv2.merge([r, g, b, new_alpha])
     
-    for filename, url in example_images.items():
-        try:
-            response = requests.get(url)
-            if response.status_code == 200:
-                with open(f"examples/{filename}", "wb") as f:
-                    f.write(response.content)
-        except Exception as e:
-            print(f"Couldn't download {filename}: {str(e)}")
-            # Provide fallback blank image
-            blank = Image.new("RGBA", (256, 256), (0, 0, 0, 0))
-            blank.save(f"examples/{filename}")
+    return Image.fromarray(result)
 
-# Download examples at startup
-download_example_images()
-
-# Create Gradio interface
-with gr.Blocks(title="✨ Edge Refiner") as demo:
+# Gradio interface
+with gr.Blocks(title="✨ Advanced Edge Refiner") as demo:
     gr.Markdown("""
-    # ✨ Edge Refiner - Clean Up Your Background-Removed Images!
-    Refine the edges of transparent PNG images for cleaner results.
+    # ✨ Advanced Edge Refiner
+    Removes leftover background artifacts around hair and fine edges
     """)
     
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="pil", label="Input Image")
-            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")
+            edge_smoothness = gr.Slider(1, 5, value=3, label="Edge Smoothness")
+            blur_radius = gr.Slider(1, 5, value=2, label="Blur Strength")
+            feather_amount = gr.Slider(0, 5, value=1, label="Feather Amount")
+            threshold = gr.Slider(0, 100, value=10, label="Edge Threshold (%)")
             submit_btn = gr.Button("Refine Edges", variant="primary")
         
         with gr.Column():
             output_image = gr.Image(type="pil", label="Refined Image")
     
-    # Use local example files
-    example_files = [f for f in os.listdir("examples") if f.endswith(".png")]
-    if example_files:
-        gr.Examples(
-            examples=[[f"examples/{f}", 3, 2, 1] for f in example_files],
-            inputs=[input_image, edge_smoothness, blur_radius, feather_amount],
-            outputs=output_image,
-            fn=refine_edges,
-            cache_examples=False
-        )
-    
     submit_btn.click(
         fn=refine_edges,
-        inputs=[input_image, edge_smoothness, blur_radius, feather_amount],
+        inputs=[input_image, edge_smoothness, blur_radius, feather_amount, threshold],
         outputs=output_image
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()