import gradio as gr
import cv2
import numpy as np
from PIL import Image

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)

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__":
    iface.launch()