From dbd4aa7798c10e3acf73cd5ee01008614a07dc7d Mon Sep 17 00:00:00 2001 From: Kader Miyanyedi Date: Fri, 18 Oct 2024 16:07:42 +0300 Subject: [PATCH] feat(ByteTrack): Allow ByteTrack to track detection without class ids --- supervision/tracker/byte_tracker/core.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index d1d567fab..f453e5237 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -276,11 +276,16 @@ def callback(frame: np.ndarray, index: int) -> np.ndarray: ``` """ + num_rows = detections.xyxy.shape[0] + class_ids = np.full(num_rows, -5) + if detections.class_id is not None: + class_ids = detections.class_id + tensors = np.hstack( ( detections.xyxy, detections.confidence[:, np.newaxis], - detections.class_id[:, np.newaxis], + class_ids[:, np.newaxis], ) ) tracks = self.update_with_tensors(tensors=tensors)