app.py

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

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

def process_image_via_api(image):
    result = client.predict(
        image=handle_file(image),
        api_name="/image"
    )
    # Convert the output tuple to PIL images and save them to temporary files
    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)
        
        # Save images to temporary files
        processed_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
        origin_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
        processed_image.save(processed_temp_path)
        origin_image.save(origin_temp_path)
        
        return (processed_temp_path, origin_temp_path)
    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 save them to temporary files
    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)
        
        # Save images to temporary files
        processed_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
        origin_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
        processed_image.save(processed_temp_path)
        origin_image.save(origin_temp_path)
        
        return (processed_temp_path, origin_temp_path)
    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="filepath")
slider1_origin = ImageSlider(label="Original Image", type="filepath")
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="filepath")
slider2_origin = ImageSlider(label="Original Image", type="filepath")
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=[chameleon], 
    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)