Update app.py

app.py

@@ -2,79 +2,68 @@ import cv2
 import numpy as np
 from PIL import Image
 import gradio as gr
-import os
 
-def refine_edges(image, edge_smoothness=3, blur_radius=2, feather_amount=1, threshold=0.1):
+def refine_edges(img, edge_aggressiveness=3, bg_removal_strength=50):
     """
-    Enhanced edge refinement that specifically targets leftover background pixels
+    Advanced edge refinement using:
+    1. Alpha matte estimation
+    2. Guided filtering
+    3. Color decontamination
     """
-    img = image.convert("RGBA")
-    np_img = np.array(img)
-    r, g, b, a = cv2.split(np_img)
+    # Convert to numpy array
+    np_img = np.array(img.convert("RGBA"))
+    rgb = np_img[..., :3].astype(np.float32) / 255.0
+    alpha = np_img[..., 3].astype(np.float32) / 255.0
     
-    # Convert parameters
-    blur_kernel = blur_radius * 2 + 1
-    smooth_iterations = edge_smoothness
+    # 1. Create trimap from alpha
+    trimap = np.zeros_like(alpha)
+    trimap[alpha > 0.95] = 1       # Definite foreground
+    trimap[alpha < 0.05] = 0       # Definite background
+    trimap[(alpha >= 0.05) & (alpha <= 0.95)] = 0.5  # Unknown area
     
-    # 1. Create a strict alpha mask (remove semi-transparent pixels)
-    _, strict_alpha = cv2.threshold(a, 254, 255, cv2.THRESH_BINARY)
+    # 2. Estimate foreground/background colors
+    fg = rgb * (trimap == 1)[..., None]
+    bg = rgb * (trimap == 0)[..., None]
     
-    # 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):
-        edge_zone = cv2.morphologyEx(edge_zone, cv2.MORPH_OPEN, kernel)
-        edge_zone = cv2.morphologyEx(edge_zone, cv2.MORPH_CLOSE, kernel)
-    
-    # 4. Apply blur only to edge zone
-    blurred_edge = cv2.GaussianBlur(edge_zone, (blur_kernel, blur_kernel), 0)
-    
-    # 5. Feathering with thresholding to remove leftover bg
-    if feather_amount > 0:
-        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
-    
-    # 6. Combine with strict alpha
-    new_alpha = cv2.bitwise_or(strict_alpha, final_edge)
+    # 3. Guided filter for alpha refinement
+    radius = edge_aggressiveness * 5
+    eps = 0.01
+    refined_alpha = cv2.ximgproc.guidedFilter(
+        guide=rgb,
+        src=alpha,
+        radius=radius,
+        eps=eps
+    )
     
-    # 7. Remove color information from transparent areas
-    r = r * (new_alpha > 0)
-    g = g * (new_alpha > 0)
-    b = b * (new_alpha > 0)
+    # 4. Color decontamination
+    bg_removal = bg_removal_strength / 100.0
+    new_rgb = (rgb - bg_removal * bg) / (1 - bg_removal * (1 - refined_alpha[..., None]))
+    new_rgb = np.clip(new_rgb, 0, 1)
     
-    # Recombine channels
-    result = cv2.merge([r, g, b, new_alpha])
+    # 5. Final alpha thresholding
+    final_alpha = np.clip(refined_alpha * 255, 0, 255).astype(np.uint8)
+    new_rgb = (new_rgb * 255).astype(np.uint8)
     
+    # Combine channels
+    result = np.concatenate([new_rgb, final_alpha[..., None]], axis=-1)
     return Image.fromarray(result)
 
-# Gradio interface
-with gr.Blocks(title="✨ Advanced Edge Refiner") as demo:
-    gr.Markdown("""
-    # ✨ Advanced Edge Refiner
-    Removes leftover background artifacts around hair and fine edges
-    """)
-    
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("## ✂️ Advanced Edge Refiner")
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(type="pil", label="Input Image")
-            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")
-        
+            img_input = gr.Image(type="pil", label="Input PNG")
+            edge_slider = gr.Slider(1, 10, value=3, label="Edge Precision")
+            bg_slider = gr.Slider(1, 100, value=50, label="BG Removal Strength")
+            process_btn = gr.Button("Refine Edges")
         with gr.Column():
-            output_image = gr.Image(type="pil", label="Refined Image")
+            img_output = gr.Image(type="pil", label="Refined Result")
     
-    submit_btn.click(
+    process_btn.click(
         fn=refine_edges,
-        inputs=[input_image, edge_smoothness, blur_radius, feather_amount, threshold],
-        outputs=output_image
+        inputs=[img_input, edge_slider, bg_slider],
+        outputs=img_output
     )
 
 if __name__ == "__main__":