Spaces:
Sleeping
Sleeping
| 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) |