import threading

import gradio as gr

import mtpnet_demo, yolop_demo


def Hex_to_RGB(hex_string):
    # your code here
    r = int(hex_string[1:3], 16)
    g = int(hex_string[3:5], 16)
    b = int(hex_string[5:7], 16)
    return [r, g, b]


class TaskThread(threading.Thread):
    def __init__(self, func, args=()):
        super(TaskThread, self).__init__()
        self.func = func
        self.args = args

    def run(self):
        self.result = self.func(*self.args)

    def getResult(self):
        try:
            return self.result
        except Exception:
            return None


def detect(path, model, task, thickness, alpha_da, alpha_ll, color1, color2, color3):
    global mtpnet, yolop
    color = [Hex_to_RGB(color1), Hex_to_RGB(color2), Hex_to_RGB(color3)]
    alpha = [alpha_da, alpha_ll]
    result, result2, result3 = None, None, None
    if 'mtpnet' in model:
        mtpnet = TaskThread(mtpnet_demo.detect, args=(path, task, thickness, color, alpha))
        mtpnet.start()

    if 'yolop' in model:
        yolop = TaskThread(yolop_demo.detect, args=(path, task, thickness, color, alpha))
        yolop.start()

    if 'mtpnet' in model:
        mtpnet.join()
        result = mtpnet.getResult()

    if 'yolop' in model:
        yolop.join()
        result2 = yolop.getResult()

    return result, result2


gr.Interface(
    fn=detect,
    inputs=
    [
        gr.Image(type='filepath', label="Input Image"),
        gr.CheckboxGroup(["mtpnet", "yolop"], value=["mtpnet", "yolop"], label="Select model"),
        gr.CheckboxGroup(["Vehicle detection", "Driving area segmentation", "Lane detection"],
                         value=["Vehicle detection", "Driving area segmentation", "Lane detection"],
                         label="Select task"),
        gr.Slider(1, 5, value=2, label="Detection box line thickness", step=1),
        gr.Slider(0.1, 1, value=0.5, label="Driving area transparency", step=0.1),
        gr.Slider(0.1, 1, value=1, label="Lane Line area transparency", step=0.1),
        gr.ColorPicker(label="Detection Box Color", value='#FFFF00'),
        gr.ColorPicker(label="Driving Area Segmentation Color", value='#00FF00'),
        gr.ColorPicker(label="Lane Line Color", value='#FF0000')
    ],
    outputs=[
        gr.Image(label="Output image by mtpnet"),
        # gr.Image(label="Output Image by yolopv2"),
        gr.Image(label="Output image by yolop")
        ],
    title="MtpNet 💪",
    examples=
    [
        ["img/1.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/12.png", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/2.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/3.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/4.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/5.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/7.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/8.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/10.jpg", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],
        ["img/11.png", ["mtpnet", "yolop"], ["Vehicle detection", "Driving area segmentation", "Lane detection"], 2, 0.5, 1, '#FFFF00', '#00FF00', '#FF0000'],

    ],
    theme='default',
    description="MtpNet 💪: demo for multi-task panoptic driving 🚗 perception network").launch(share=False)