import gradio as gr
from rembg import remove
import numpy as np

# Функция для удаления водяных знаков с изображения
def remove_watermark(input_image):
    # Сохранение массива numpy в файл
    input_image = input_image.astype(np.uint8)
    input_image = np.squeeze(input_image)
    with open("input.png", "wb") as f:
        f.write(input_image)

    # Загрузка изображения и удаление фона
    output = remove("input.png")

    return output.tobytes()

# Создание Gradio приложения
iface = gr.Interface(
    fn=remove_watermark,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(),
)

iface.launch()