"""
    reference: https://github.com/xuebinqin/DIS
"""

import os

import gdown
import gradio as gr

from DIS.IsNetPipeLine import IsNetPipeLine

save_model_path = "DIS/save_models"
model_name = os.path.join(save_model_path, "isnet.pth")
# Download official weights
if not os.path.exists(model_name):
    if not os.path.exists(save_model_path):
        os.mkdir(save_model_path)
    MODEL_PATH_URL = "https://huggingface.co/Superlang/ImageProcess/resolve/main/isnet.pth"
    gdown.download(MODEL_PATH_URL, model_name, use_cookies=False)

pipe = IsNetPipeLine(model_path=model_name)


def inference(image):
    return pipe(image)


title = "remove background"
interface = gr.Interface(
    fn=inference,
    inputs=gr.Image(type='pil'),
    outputs=["image", "image"],
    title=title,
    allow_flagging='never',
    cache_examples=True,
).queue(concurrency_count=1, api_open=True).launch(show_api=True, show_error=True)