import os
os.system("pip3 install cython_bbox gdown 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'")
from torchyolo import YoloHub
import gradio as gr
from utils import attempt_download_from_hub

""" 
Paper Implementation
#"kadirnar/OcSort"
#"kadirnar/SORT"
#"kadirnar/ByteTracker"
#"kadirnar/strongsort"
"""

def object_tracker(
    source: str,
    model_type: str,
    model_path: str,
    tracker_type: str,
    tracker_config_path: str,
    StrongSort_OsNet_Path: str = None,
):
    model = YoloHub(
        config_path="default_config.yaml",
        model_type=model_type,
        model_path=model_path,
    )

    if tracker_type == "STRONGSORT":
        StrongSort_OsNet_Path = attempt_download_from_hub(StrongSort_OsNet_Path)

    model.predict(
        source=source,
        tracker_type=tracker_type,
        tracker_weight_path=StrongSort_OsNet_Path,
        tracker_config_path=tracker_config_path,
    )

    return 'output.mp4'
       

inputs = [
    gr.Image(),
    gr.inputs.Dropdown(
        label="Model Type",
        choices=["yolov5", "yolov6", "yolov8"],
        default="yolov5",
    ),
    gr.inputs.Dropdown(
        label="Model Path",
        choices=[
            "kadirnar/yolov5s6-v6.0",
            "kadirnar/yolov6m-v3.0",
            "kadirnar/yolov8n-v8.0",
        ],
        default="kadirnar/yolov5s6-v6.",
    ),
    gr.inputs.Dropdown(
        label="Tracker Type",
        choices=["NORFAIR", "STRONGSORT", "OCSORT", "BYTETRACK", "SORT"],
        default="NORFAIR",
    ),
    gr.inputs.Dropdown(
        label="Tracker Config Path",
        choices=[
            "tracker/norfair_track.yaml",
            "tracker/strong_sort.yaml",
            "tracker/oc_sort.yaml",
            "tracker/byte_track.yaml",
            "tracker/sort_track.yaml",
        ],
        default="tracker/norfair_track.yaml",
    ),
    gr.inputs.Dropdown(
        label="Tracker Weight Path",
        choices=[
            "kadirnar/osnet_x0_5_imagenet",
            "kadirnar/osnet_x1_0_imagenet",
            "kadirnar/osnet_x0_25_imagenet"
        ],
        default="kadirnar/osnet_x0_5_imagenet",
    ),
]
examples = [
    [
        "test.mp4",
        "yolov5",
        "kadirnar/yolov5s6-v6.0",
        "SORT",
        "tracker/sort_track.yaml",
    ],
    [
        "testv2.mp4",
        "yolov6",
        "kadirnar/yolov6m-v3.0",
        "OCSORT",
        "tracker/oc_sort.yaml"
    ]
]
outputs = gr.Image()
title = "TorchYolo: YOLO Series Object Detection and Track Algorithm Library"

demo_app = gr.Interface(
    fn=object_tracker,
    inputs=inputs,
    examples=examples,
    outputs=outputs,
    title=title,
    cache_examples=False,
    theme='huggingface',
)
demo_app.launch(debug=True, enable_queue=True)