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

Fix error with torch tensor input in model.track() (#19278)

Browse source 
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Mohammed Yasin 2025-02-18 19:59:30 +08:00 • committed by GitHub
parent e981fc629d
commit 0f81777af5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 5 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

12
ultralytics/trackers/track.py
View file

@ -66,25 +66,23 @@ def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None
>>> predictor = YourPredictorClass() >>> predictor = YourPredictorClass()
>>> on_predict_postprocess_end(predictor, persist=True) >>> on_predict_postprocess_end(predictor, persist=True)
""" """
path, im0s = predictor.batch[:2]
is_obb = predictor.args.task == "obb" is_obb = predictor.args.task == "obb"
is_stream = predictor.dataset.mode == "stream" is_stream = predictor.dataset.mode == "stream"
for i in range(len(im0s)): for i, result in enumerate(predictor.results):
tracker = predictor.trackers[i if is_stream else 0] tracker = predictor.trackers[i if is_stream else 0]
vid_path = predictor.save_dir / Path(path[i]).name vid_path = predictor.save_dir / Path(result.path).name
if not persist and predictor.vid_path[i if is_stream else 0] != vid_path: if not persist and predictor.vid_path[i if is_stream else 0] != vid_path:
tracker.reset() tracker.reset()
predictor.vid_path[i if is_stream else 0] = vid_path 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() det = (result.obb if is_obb else result.boxes).cpu().numpy()
if len(det) == 0: if len(det) == 0:
continue continue
tracks = tracker.update(det, im0s[i]) tracks = tracker.update(det, result.orig_img)
if len(tracks) == 0: if len(tracks) == 0:
continue continue
idx = tracks[:, -1].astype(int) idx = tracks[:, -1].astype(int)
predictor.results[i] = predictor.results[i][idx] predictor.results[i] = result[idx]
update_args = {"obb" if is_obb else "boxes": torch.as_tensor(tracks[:, :-1])} update_args = {"obb" if is_obb else "boxes": torch.as_tensor(tracks[:, :-1])}
predictor.results[i].update(**update_args) predictor.results[i].update(**update_args)