import gradio as gr
from gradio_client import Client, handle_file
from gradio_imageslider import ImageSlider
from PIL import Image
import tempfile
import os

# Инициализируем клиент
client = Client("not-lain/background-removal")

def process_image_via_api(image):
    # Сохраняем изображение во временный файл
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
        image.save(temp_file.name)
        temp_file_path = temp_file.name
    
    result = client.predict(
        image=handle_file(temp_file_path),
        api_name="/image"
    )
    
    # Удаляем временный файл после использования
    os.remove(temp_file_path)
    
    # Convert the output tuple to PIL images and return
    if result:
        processed_image_path = result[0]
        origin_image_path = result[1]
        processed_image = Image.open(processed_image_path)
        origin_image = Image.open(origin_image_path)
        return (processed_image, origin_image)
    return None, None

def process_url_via_api(url):
    result = client.predict(
        image=url,
        api_name="/text"
    )
    # Convert the output tuple to PIL images and return
    if result:
        processed_image_path = result[0]
        origin_image_path = result[1]
        processed_image = Image.open(processed_image_path)
        origin_image = Image.open(origin_image_path)
        return (processed_image, origin_image)
    return None, None

def process_file_via_api(f):
    result = client.predict(
        f=handle_file(f),
        api_name="/png"
    )
    # Return the path to the saved PNG file
    if result:
        return result
    return None

# Пример изображений
chameleon = "butterfly.jpg"
url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"

# Tab 1: Image Upload
slider1_processed = ImageSlider(label="Processed Image", type="pil")
slider1_origin = ImageSlider(label="Original Image", type="pil")
image_upload = gr.Image(label="Upload an image")
tab1 = gr.Interface(
    fn=process_image_via_api, 
    inputs=image_upload, 
    outputs=[slider1_processed, slider1_origin], 
    examples=[chameleon], 
    api_name="/image_api"
)

# Tab 2: URL Input
slider2_processed = ImageSlider(label="Processed Image", type="pil")
slider2_origin = ImageSlider(label="Original Image", type="pil")
url_input = gr.Textbox(label="Paste an image URL")
tab2 = gr.Interface(
    fn=process_url_via_api, 
    inputs=url_input, 
    outputs=[slider2_processed, slider2_origin], 
    examples=[url_example], 
    api_name="/url_api"
)

# Tab 3: File Output
output_file = gr.File(label="Output PNG File")
image_file_upload = gr.Image(label="Upload an image", type="filepath")
tab3 = gr.Interface(
    fn=process_file_via_api, 
    inputs=image_file_upload, 
    outputs=output_file, 
    examples=["butterfly.jpg"], 
    api_name="/png_api"
)

# Создаем интерфейс с вкладками
demo = gr.TabbedInterface(
    [tab1, tab2, tab3], 
    ["Image Upload", "URL Input", "File Output"], 
    title="Background Removal Tool"
)

if __name__ == "__main__":
    demo.launch(show_error=True)