import track
import os
from pathlib import Path

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # yolov5 strongsort root directory
WEIGHTS = ROOT / 'weights'
trackers =  "strongsort" # ["strongsort", "bytetrack"]
path = "./video"
v = [x for x in os.listdir(path) if x.endswith(".mp4")]

for i in v:
    # , classes=[0, 1, 2, 3, 5, 7, 9, 11, 10]
    track.run(source=ROOT / f"video/{i}", tracking_method=trackers, yolo_weights=WEIGHTS / "best2.pt", project=ROOT / f"runs/{trackers}3", name=f"video{i.split('.')[0]}", iou_thres=0.45, conf_thres=0.35, hide_conf=True)
    print(f"{i}检测完毕")