import os
import numpy as np
import onnxruntime as ort
from PIL import Image
import gradio as gr
import torchvision.transforms as transforms
import time

# Выбираем доступный execution provider
available_providers = ort.get_available_providers()
execution_provider = "CUDAExecutionProvider" if "CUDAExecutionProvider" in available_providers else "CPUExecutionProvider"

# Загружаем модель ONNX
onnx_session = ort.InferenceSession("BEN2_Base.onnx", providers=[execution_provider])

# Трансформации
transform = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
])

# Удаление фона с замером времени
def remove_background(image: Image.Image):
    start = time.time()
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    orig_size = image.size
    image = image.convert("RGB")
    img_tensor = transform(image).unsqueeze(0).numpy()

    ort_inputs = {onnx_session.get_inputs()[0].name: img_tensor}
    ort_outs = onnx_session.run(None, ort_inputs)
    mask = ort_outs[0][0][0]

    matte = (mask * 255).astype(np.uint8)
    matte_img = Image.fromarray(matte).resize(orig_size, Image.BILINEAR)

    image = image.resize(orig_size).convert("RGBA")
    image.putalpha(matte_img)

    elapsed = round(time.time() - start, 2)
    return image, f"Estimated processing time: {elapsed} seconds"

# Gradio интерфейс
with gr.Blocks() as demo:
    gr.Markdown("# BEN2 Background Remover")
    with gr.Row():
        input_image = gr.Image(label="Upload Image")
        output_image = gr.Image(label="Processed Image")
    with gr.Row():
        eta = gr.Textbox(label="Processing Time (ETA)")
    with gr.Row():
        submit = gr.Button("Remove Background")
        clear = gr.Button("Clear")

    submit.click(fn=remove_background, inputs=input_image, outputs=[output_image, eta])
    clear.click(lambda: (None, None, ""), outputs=[input_image, output_image, eta])

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