You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
80 lines
2.9 KiB
80 lines
2.9 KiB
from functools import partial |
|
from pathlib import Path |
|
|
|
import torch |
|
|
|
from ultralytics.utils import IterableSimpleNamespace, yaml_load |
|
from ultralytics.utils.checks import check_yaml |
|
|
|
from .bot_sort import BOTSORT |
|
from .byte_tracker import BYTETracker |
|
|
|
# A mapping of tracker types to corresponding tracker classes |
|
TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT} |
|
|
|
|
|
def on_predict_start(predictor: object, persist: bool = False) -> None: |
|
""" |
|
Initialize trackers for object tracking during prediction. |
|
|
|
Args: |
|
predictor (object): The predictor object to initialize trackers for. |
|
persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False. |
|
|
|
Raises: |
|
AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'. |
|
""" |
|
if predictor.args.task == 'obb': |
|
raise NotImplementedError('ERROR ❌ OBB task does not support track mode!') |
|
if hasattr(predictor, 'trackers') and persist: |
|
return |
|
|
|
tracker = check_yaml(predictor.args.tracker) |
|
cfg = IterableSimpleNamespace(**yaml_load(tracker)) |
|
|
|
if cfg.tracker_type not in ['bytetrack', 'botsort']: |
|
raise AssertionError(f"Only 'bytetrack' and 'botsort' are supported for now, but got '{cfg.tracker_type}'") |
|
|
|
trackers = [] |
|
for _ in range(predictor.dataset.bs): |
|
tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30) |
|
trackers.append(tracker) |
|
predictor.trackers = trackers |
|
|
|
|
|
def on_predict_postprocess_end(predictor: object, persist: bool = False) -> None: |
|
""" |
|
Postprocess detected boxes and update with object tracking. |
|
|
|
Args: |
|
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] |
|
|
|
for i in range(bs): |
|
if not persist and predictor.vid_path[i] != str(predictor.save_dir / Path(path[i]).name): # new video |
|
predictor.trackers[i].reset() |
|
|
|
det = predictor.results[i].boxes.cpu().numpy() |
|
if len(det) == 0: |
|
continue |
|
tracks = predictor.trackers[i].update(det, im0s[i]) |
|
if len(tracks) == 0: |
|
continue |
|
idx = tracks[:, -1].astype(int) |
|
predictor.results[i] = predictor.results[i][idx] |
|
predictor.results[i].update(boxes=torch.as_tensor(tracks[:, :-1])) |
|
|
|
|
|
def register_tracker(model: object, persist: bool) -> None: |
|
""" |
|
Register tracking callbacks to the model for object tracking during prediction. |
|
|
|
Args: |
|
model (object): The model object to register tracking callbacks for. |
|
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_postprocess_end', partial(on_predict_postprocess_end, persist=persist))
|
|
|