import gradio as gr
import torch
from carvekit.api.interface import Interface
from carvekit.ml.wrap.basnet import BASNET
from carvekit.ml.wrap.deeplab_v3 import DeepLabV3
from carvekit.ml.wrap.fba_matting import FBAMatting
from carvekit.ml.wrap.tracer_b7 import TracerUniversalB7
from carvekit.ml.wrap.u2net import U2NET
from carvekit.pipelines.postprocessing import MattingMethod
from carvekit.pipelines.preprocessing import PreprocessingStub
from carvekit.trimap.generator import TrimapGenerator

device = 'cuda' if torch.cuda.is_available() else 'cpu'

segment_net = {
    "U2NET": U2NET(device=device, batch_size=1),
    "BASNET": BASNET(device=device, batch_size=1),
    "DeepLabV3": DeepLabV3(device=device, batch_size=1),
    "TracerUniversalB7": TracerUniversalB7(device=device, batch_size=1)
}

fba = FBAMatting(device=device,
                 input_tensor_size=2048,
                 batch_size=1)

trimap = TrimapGenerator()

preprocessing = PreprocessingStub()

postprocessing = MattingMethod(matting_module=fba,
                               trimap_generator=trimap,
                               device=device)

method_choices = [k for k, v in segment_net.items()]


def generate_trimap(method, original):
    mask = segment_net[method]([original])
    return trimap(original_image=original, mask=mask[0])


def predict(method, image):
    method = segment_net[method]
    return Interface(pre_pipe=preprocessing,
                     post_pipe=postprocessing,
                     seg_pipe=method)([image])[0]


footer = r"""
<center>
<img src='https://raw.githubusercontent.com/leonelhs/image-background-remove-tool/master/docs/imgs/logo.png' alt='CarveKit' width="200" height="80">
</br>
<b>
Demo based on <a href='https://github.com/OPHoperHPO/image-background-remove-tool'>CarveKit</a>
</b>
</center>
"""

with gr.Blocks(title="CarveKit") as app:
    gr.Markdown("<center><h1><b>CarveKit</b></h1></center>")
    gr.HTML("<center><h3>High-quality image background removal</h3></center>")

    with gr.Tabs() as tabs:
        with gr.TabItem("Remove background", id=0):
            with gr.Row(equal_height=False):
                with gr.Column():
                    input_img = gr.Image(type="pil", label="Input image")
                    drp_itf = gr.Dropdown(
                        value="TracerUniversalB7",
                        label="Segmentor model",
                        choices=method_choices)
                    run_btn = gr.Button(variant="primary")
                with gr.Column():
                    output_img = gr.Image(type="pil", label="result")

            run_btn.click(predict, [drp_itf, input_img], [output_img])

        with gr.TabItem("Trimap generator", id=1):
            with gr.Row(equal_height=False):
                with gr.Column():
                    trimap_input = gr.Image(type="pil", label="Input image")
                    drp_itf = gr.Dropdown(
                        value="TracerUniversalB7",
                        label="Segmentor model",
                        choices=method_choices)
                    trimap_btn = gr.Button(variant="primary")
                with gr.Column():
                    trimap_output = gr.Image(type="pil", label="result")

            trimap_btn.click(generate_trimap, [drp_itf, trimap_input], [trimap_output])

    with gr.Row():
        gr.HTML(footer)

app.launch(share=False, debug=True, enable_queue=True, show_error=True)