Update app.py

app.py

@@ -1,70 +1,39 @@
+import gradio as gr
 import cv2
 import numpy as np
 from PIL import Image
-import gradio as gr
 
-def refine_edges(img, edge_aggressiveness=3, bg_removal_strength=50):
-    """
-    Advanced edge refinement using:
-    1. Alpha matte estimation
-    2. Guided filtering
-    3. Color decontamination
-    """
-    # 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
-    
-    # 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
-    
-    # 2. Estimate foreground/background colors
-    fg = rgb * (trimap == 1)[..., None]
-    bg = rgb * (trimap == 0)[..., None]
-    
-    # 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
-    )
-    
-    # 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)
-    
-    # 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)
+def refine_edges(image: Image.Image):
+    # Convert PIL Image to NumPy array
+    image_np = np.array(image)
+
+    # Check if image has alpha channel
+    if image_np.shape[2] != 4:
+        return image  # Return original if no alpha channel
+
+    # Separate color and alpha channels
+    b, g, r, a = cv2.split(image_np)
+
+    # Apply Gaussian blur to alpha channel
+    a_blurred = cv2.GaussianBlur(a, (5, 5), 0)
+
+    # Apply morphological operations to smooth edges
+    kernel = np.ones((3, 3), np.uint8)
+    a_morph = cv2.morphologyEx(a_blurred, cv2.MORPH_OPEN, kernel)
+
+    # Merge channels back
+    result = cv2.merge((b, g, r, a_morph))
+
+    # Convert back to PIL Image
     return Image.fromarray(result)
 
-# Gradio Interface
-with gr.Blocks() as demo:
-    gr.Markdown("## ✂️ Advanced Edge Refiner")
-    with gr.Row():
-        with gr.Column():
-            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():
-            img_output = gr.Image(type="pil", label="Refined Result")
-    
-    process_btn.click(
-        fn=refine_edges,
-        inputs=[img_input, edge_slider, bg_slider],
-        outputs=img_output
-    )
+iface = gr.Interface(
+    fn=refine_edges,
+    inputs=gr.Image(type="pil", image_mode="RGBA", label="Upload Transparent PNG"),
+    outputs=gr.Image(type="pil", image_mode="RGBA", label="Refined Image"),
+    title="Edge Smoother",
+    description="Upload a transparent PNG image to refine its edges, especially around hair and fur."
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()