import spaces
import gradio as gr
from transparent_background import Remover
from PIL import Image
import numpy as np
import io
import tempfile

@spaces.GPU
def remove_background(image):
    remover = Remover()
    if isinstance(image, Image.Image):
        output = remover.process(image)
    elif isinstance(image, np.ndarray):
        image_pil = Image.fromarray(image)
        output = remover.process(image_pil)
    else:
        raise TypeError("Unsupported image type")
    
    # 保存PNG图片到临时文件
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    output.save(temp_file, format="PNG")
    temp_file.close()

    return temp_file.name

# 界面设置
iface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(label="Upload Image"),
    outputs=gr.Image(label="Output Image", type="filepath"),
    title="AI Background Remover - Automatically Remove Image Backgrounds",
    description="Upload an image and our AI-powered background remover will automatically make the background transparent. Perfect for various needs requiring transparent images.",
    article="""
    ## How to Use
    1. Click **Upload Image** or drag and drop an image into the upload area.
    2. Click the **Submit** button and wait for a moment. The processed image will appear in the output area on the right.
    3. The background of the processed image will be transparent, and the image will be in PNG format for easy use.

    ### Application Scenarios
    - **E-commerce**: Remove backgrounds from product images to create clean product photos.
    - **Graphic Design**: Quickly get images with transparent backgrounds for composition and design.
    - **Social Media**: Create avatars or icons with transparent backgrounds that blend perfectly on various backgrounds.

    Our tool leverages the latest AI technology to efficiently and accurately remove image backgrounds, saving you time and effort.
    """
)

if __name__ == "__main__":
    iface.launch()