HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.1.27 batched tracking fixes (#8842)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2024-03-12 02:29:41 +08:00 • committed by GitHub
parent 3555785167
commit 2ea6b2b889
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 38 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

15
ultralytics/trackers/track.py
View file

@ -38,6 +38,8 @@ def on_predict_start(predictor: object, persist: bool = False) -> None:
for _ in range(predictor.dataset.bs):
tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
trackers.append(tracker)
if predictor.dataset.mode != "stream": # only need one tracker for other modes.
break
predictor.trackers = trackers
predictor.vid_path = [None] * predictor.dataset.bs # for determining when to reset tracker on new video
@ -50,20 +52,21 @@ def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None
predictor (object): The predictor object containing the predictions.
persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False.
"""
bs = predictor.dataset.bs
path, im0s = predictor.batch[:2]
is_obb = predictor.args.task == "obb"
for i in range(bs):
is_stream = predictor.dataset.mode == "stream"
for i in range(len(im0s)):
tracker = predictor.trackers[i if is_stream else 0]
vid_path = predictor.save_dir / Path(path[i]).name
if not persist and predictor.vid_path[i] != vid_path: # new video
predictor.trackers[i].reset()
predictor.vid_path[i] = vid_path
if not persist and predictor.vid_path[i if is_stream else 0] != vid_path:
tracker.reset()
predictor.vid_path[i if is_stream else 0] = vid_path
det = (predictor.results[i].obb if is_obb else predictor.results[i].boxes).cpu().numpy()
if len(det) == 0:
continue
tracks = predictor.trackers[i].update(det, im0s[i])
tracks = tracker.update(det, im0s[i])
if len(tracks) == 0:
continue
idx = tracks[:, -1].astype(int)