import os

import gradio as gr

from inference import inference

input_video = gr.Video(mirror_webcam=False)

dd_model = gr.Dropdown(choices=["YOLOv7", "YOLOv7 Tiny"], value="YOLOv7", label="Model")

features = gr.CheckboxGroup(
    choices=["Track camera movement", "Draw objects paths"],
    value=["Track camera movement", "Draw objects paths"],
    label="Features",
    type="index",
)

cb_path_draw = gr.Checkbox(value=True, label="Draw objects paths")

dd_track_points = gr.Dropdown(
    choices=["Bounding box", "Centroid"], value="Bounding box", label="Detections style"
)

slide_threshold = gr.Slider(minimum=0, maximum=1, value=0.25, label="Model confidence threshold")

intput_components = [
        input_video,
        dd_model,
        features,
        dd_track_points,
        slide_threshold
    ]

output_components = "playablevideo"

example_list = [["examples/" + example] for example in os.listdir("examples")]

iface = gr.Interface(
    fn=inference,
    inputs=intput_components,
    outputs=output_components,
    examples=example_list,
    cache_examples=True,
)

iface.launch()