ultralytics 8.0.207 fix model.track(persist=True) bug (#6145)

This commit is contained in:
Glenn Jocher 2023-11-05 15:22:00 +01:00 committed by GitHub
parent 168e536fae
commit 70c400ee15
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 4 deletions

View File

@ -1,6 +1,6 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license # Ultralytics YOLO 🚀, AGPL-3.0 license
__version__ = '8.0.206' __version__ = '8.0.207'
from ultralytics.models import RTDETR, SAM, YOLO from ultralytics.models import RTDETR, SAM, YOLO
from ultralytics.models.fastsam import FastSAM from ultralytics.models.fastsam import FastSAM

View File

@ -38,13 +38,13 @@ def on_predict_start(predictor, persist=False):
predictor.trackers = trackers predictor.trackers = trackers
def on_predict_postprocess_end(predictor): def on_predict_postprocess_end(predictor, persist=False):
"""Postprocess detected boxes and update with object tracking.""" """Postprocess detected boxes and update with object tracking."""
bs = predictor.dataset.bs bs = predictor.dataset.bs
path, im0s = predictor.batch[:2] path, im0s = predictor.batch[:2]
for i in range(bs): for i in range(bs):
if predictor.vid_path[i] != str(predictor.save_dir / Path(path[i]).name): # new video if not persist and predictor.vid_path[i] != str(predictor.save_dir / Path(path[i]).name): # new video
predictor.trackers[i].reset() predictor.trackers[i].reset()
det = predictor.results[i].boxes.cpu().numpy() det = predictor.results[i].boxes.cpu().numpy()
@ -67,4 +67,4 @@ def register_tracker(model, persist):
persist (bool): Whether to persist the trackers if they already exist. persist (bool): Whether to persist the trackers if they already exist.
""" """
model.add_callback('on_predict_start', partial(on_predict_start, persist=persist)) model.add_callback('on_predict_start', partial(on_predict_start, persist=persist))
model.add_callback('on_predict_postprocess_end', on_predict_postprocess_end) model.add_callback('on_predict_postprocess_end', partial(on_predict_postprocess_end, persist=persist))