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

def remove_background(image):
    try:
        # PIL → bytes
        img = Image.fromarray(np.array(image))

        # 🔥 AI Background Remove
        output = remove(img)

        # RGBA image
        output = np.array(output)

        # 🔥 Edge smooth (important)
        if output.shape[2] == 4:
            alpha = output[:, :, 3]

            # smooth mask
            alpha = cv2.GaussianBlur(alpha, (5,5), 0)

            output[:, :, 3] = alpha

        return Image.fromarray(output)

    except Exception as e:
        print("ERROR:", e)
        return image


demo = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="🔥 AI Background Remover (Pro)",
    description="Clean Cut + Smooth Hair + HD Output"
)

demo.launch()