Spaces:
Sleeping
Sleeping
File size: 2,771 Bytes
5265c21 3b5e064 c1786f2 8005a73 c1786f2 39c7e92 8005a73 c1786f2 8005a73 c1786f2 8005a73 c1786f2 8005a73 c1786f2 8005a73 c1786f2 2be2cf6 c1786f2 8005a73 c1786f2 15bbe04 c1786f2 15bbe04 c1786f2 9da273f c1786f2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
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) |