yolov5_tracking / test.py
xfys's picture
Upload 645 files
47af768
raw
history blame
606 Bytes
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}检测完毕")