import gradio as gr
from ultralytics import YOLO

#各モデルをロード
detect_model = YOLO("best.pt")
seg_model = YOLO('yolov8n-seg.pt')
cls_model = YOLO('yolov8n-cls.pt')
pose_model = YOLO("yolov8n-pose.pt")

#オプションリストをオブジェクトに整形
def return_options(checkbox):
    option = {}

    #整形
    for check in checkbox:
        option[check] = True

    return option

#結果を描画する
def plot(res):
    plot = res[0].plot()
    return plot

def inference(type,input,conf):  
    if(type == "Detección"):
        model = detect_model
    elif(type == "seg"):
        model = seg_model
    elif(type == "cls"):
        model = cls_model
    elif(type == "Pose"):
        model = pose_model

    cpu=True
    if(cpu):
        device = "cpu"

    
    line_width = None

    
    
    #物体検出を実行
    res = model(input,conf=conf,iou=0.7,device="cpu",max_det=300,line_width=line_width)

    #結果を描画
    plotted = plot(res)
    return plotted

with gr.Blocks() as app:
    #ヘッダー
    gr.Markdown("# Detección de comportamiento sospechoso")

    #タブ
    with gr.Tabs():
        #inferenceタブ
        with gr.TabItem("Inferencia"):
            with gr.Row():
                input = gr.Image()
                output = gr.Image()

            type = gr.Radio(["Detección", "Pose"], value="Detección", label="Tasks")

            #オプション 
            conf = gr.Slider(minimum=0, maximum=1, value=0.25, step=0.01, interactive=True,label="conf")
            #iou = gr.Slider(minimum=0, maximum=1, value=0.7, step=0.01, interactive=True,label="iou")
            #checkbox = gr.CheckboxGroup(["half","show","save","save_txt","save_conf","save_crop","hide_labels","hide_conf","vid_stride","visualize","augment","agnostic_nms","retina_masks","boxes"], label="Options",value=["boxes"])

            #device = gr.Number(value=0, label="device", interactive=True, precision=0)
            #cpu = gr.Checkbox(label="cpu", interactive=True)
            #max_det = gr.Number(value=300, label="max_det", interactive=True, precision=0)
            #line_width = gr.Number(value=0, label="line_width", interactive=True, precision=0)

            inference_button = gr.Button("Realizar Inferencia")



    #inputから画像を取得してdetect関数を実行
    inference_button.click(inference, inputs=[type,input,conf], outputs=output)