Skip to content

Reference for ultralytics/trackers/track.py

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/trackers/track.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!


ultralytics.trackers.track.on_predict_start

on_predict_start(predictor: object, persist: bool = False) -> None

Initialize trackers for object tracking during prediction.

Parameters:

Name Type Description Default
predictor object

The predictor object to initialize trackers for.

required
persist bool

Whether to persist the trackers if they already exist.

False

Raises:

Type Description
AssertionError

If the tracker_type is not 'bytetrack' or 'botsort'.

ValueError

If the task is 'classify' as classification doesn't support tracking.

Examples:

Initialize trackers for a predictor object:

>>> predictor = SomePredictorClass()
>>> on_predict_start(predictor, persist=True)
Source code in ultralytics/trackers/track.py
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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): Whether to persist the trackers if they already exist.

    Raises:
        AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'.
        ValueError: If the task is 'classify' as classification doesn't support tracking.

    Examples:
        Initialize trackers for a predictor object:
        >>> predictor = SomePredictorClass()
        >>> on_predict_start(predictor, persist=True)
    """
    if predictor.args.task == "classify":
        raise ValueError("❌ Classification doesn't support 'mode=track'")

    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}'")

    if cfg.tracker_type == "botsort" and cfg.with_reid and cfg.model == "auto":
        from ultralytics.nn.modules.head import Detect

        if not (
            isinstance(predictor.model.model, torch.nn.Module)
            and isinstance(predictor.model.model.model[-1], Detect)
            and not predictor.model.model.model[-1].end2end
        ):
            cfg.model = "yolo11n-cls.pt"
        else:
            predictor.save_feats = True
            predictor._feats = None

            # Register hook to extract input of Detect layer
            def capture_io(module, input, output):
                predictor._feats = input[0]

            predictor.model.model.model[-1].register_forward_hook(capture_io)

    trackers = []
    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





ultralytics.trackers.track.on_predict_postprocess_end

on_predict_postprocess_end(predictor: object, persist: bool = False) -> None

Postprocess detected boxes and update with object tracking.

Parameters:

Name Type Description Default
predictor object

The predictor object containing the predictions.

required
persist bool

Whether to persist the trackers if they already exist.

False

Examples:

Postprocess predictions and update with tracking

>>> predictor = YourPredictorClass()
>>> on_predict_postprocess_end(predictor, persist=True)
Source code in ultralytics/trackers/track.py
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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): Whether to persist the trackers if they already exist.

    Examples:
        Postprocess predictions and update with tracking
        >>> predictor = YourPredictorClass()
        >>> on_predict_postprocess_end(predictor, persist=True)
    """
    is_obb = predictor.args.task == "obb"
    is_stream = predictor.dataset.mode == "stream"
    for i, result in enumerate(predictor.results):
        tracker = predictor.trackers[i if is_stream else 0]
        vid_path = predictor.save_dir / Path(result.path).name
        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 = (result.obb if is_obb else result.boxes).cpu().numpy()
        if len(det) == 0:
            continue
        tracks = tracker.update(det, result.orig_img, getattr(result, "feats", None))
        if len(tracks) == 0:
            continue
        idx = tracks[:, -1].astype(int)
        predictor.results[i] = result[idx]

        update_args = {"obb" if is_obb else "boxes": torch.as_tensor(tracks[:, :-1])}
        predictor.results[i].update(**update_args)





ultralytics.trackers.track.register_tracker

register_tracker(model: object, persist: bool) -> None

Register tracking callbacks to the model for object tracking during prediction.

Parameters:

Name Type Description Default
model object

The model object to register tracking callbacks for.

required
persist bool

Whether to persist the trackers if they already exist.

required

Examples:

Register tracking callbacks to a YOLO model

>>> model = YOLOModel()
>>> register_tracker(model, persist=True)
Source code in ultralytics/trackers/track.py
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
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.

    Examples:
        Register tracking callbacks to a YOLO model
        >>> model = YOLOModel()
        >>> register_tracker(model, persist=True)
    """
    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))





📅 Created 1 year ago ✏️ Updated 7 months ago