From c0b87a6f042a11a2713bc58f60e83cfac2e4ae72 Mon Sep 17 00:00:00 2001 From: Antonio Consiglio <77753494+AntonioConsiglio@users.noreply.github.com> Date: Sat, 7 Oct 2023 19:03:30 +0200 Subject: [PATCH 1/6] Make sv.ByteTrack work with segmentation model#418 --- supervision/tracker/byte_tracker/core.py | 30 +++++++++++++++++------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index 34a90138..8ddf0963 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -11,7 +11,7 @@ from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter class STrack(BaseTrack): shared_kalman = KalmanFilter() - def __init__(self, tlwh, score, class_ids): + def __init__(self, tlwh, score, class_ids, mask=None): # wait activate self._tlwh = np.asarray(tlwh, dtype=np.float32) self.kalman_filter = None @@ -21,6 +21,7 @@ class STrack(BaseTrack): self.score = score self.class_ids = class_ids self.tracklet_len = 0 + self.mask = mask def predict(self): mean_state = self.mean.copy() @@ -74,6 +75,7 @@ class STrack(BaseTrack): if new_id: self.track_id = self.next_id() self.score = new_track.score + self.mask = new_track.mask if new_track.mask is not None else None def update(self, new_track, frame_id): """ @@ -95,6 +97,8 @@ class STrack(BaseTrack): self.score = new_track.score + self.mask = new_track.mask if new_track.mask is not None else None + @property def tlwh(self): """Get current position in bounding box format `(top left x, top left y, @@ -231,9 +235,9 @@ class ByteTrack: ... ) ``` """ - tracks = self.update_with_tensors( - tensors=detections2boxes(detections=detections) + tensors=detections2boxes(detections=detections), + masks = detections.mask ) detections = Detections.empty() if len(tracks) > 0: @@ -249,12 +253,16 @@ class ByteTrack: detections.confidence = np.array( [t.score for t in tracks], dtype=np.float32 ) + detections.mask = np.array( + [t.mask for t in tracks], dtype=bool + ) + else: detections.tracker_id = np.array([], dtype=int) return detections - def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]: + def update_with_tensors(self, tensors: np.ndarray, masks:np.ndarray = None) -> List[STrack]: """ Updates the tracker with the provided tensors and returns the updated tracks. @@ -281,6 +289,12 @@ class ByteTrack: inds_second = np.logical_and(inds_low, inds_high) dets_second = bboxes[inds_second] dets = bboxes[remain_inds] + if masks is not None: + masks_hs = masks[remain_inds] + masks_ls = masks[inds_second] + else: + masks_hs = np.array([None]*len(remain_inds)) + masks_ls = np.array([None]*len(inds_second)) scores_keep = scores[remain_inds] scores_second = scores[inds_second] @@ -290,8 +304,8 @@ class ByteTrack: if len(dets) > 0: """Detections""" detections = [ - STrack(STrack.tlbr_to_tlwh(tlbr), s, c) - for (tlbr, s, c) in zip(dets, scores_keep, class_ids_keep) + STrack(STrack.tlbr_to_tlwh(tlbr), s, c,m) + for (tlbr, s, c,m) in zip(dets, scores_keep, class_ids_keep,masks_hs) ] else: detections = [] @@ -331,8 +345,8 @@ class ByteTrack: if len(dets_second) > 0: """Detections""" detections_second = [ - STrack(STrack.tlbr_to_tlwh(tlbr), s, c) - for (tlbr, s, c) in zip(dets_second, scores_second, class_ids_second) + STrack(STrack.tlbr_to_tlwh(tlbr), s, c,m) + for (tlbr, s, c, m) in zip(dets_second, scores_second, class_ids_second,masks_ls) ] else: detections_second = [] From 99707582bc75a1b7077c2090a0ef717ce54204f2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Oct 2023 17:08:16 +0000 Subject: [PATCH 2/6] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/tracker/byte_tracker/core.py | 25 ++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index 8ddf0963..c7d6ae18 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -236,8 +236,7 @@ class ByteTrack: ``` """ tracks = self.update_with_tensors( - tensors=detections2boxes(detections=detections), - masks = detections.mask + tensors=detections2boxes(detections=detections), masks=detections.mask ) detections = Detections.empty() if len(tracks) > 0: @@ -253,16 +252,16 @@ class ByteTrack: detections.confidence = np.array( [t.score for t in tracks], dtype=np.float32 ) - detections.mask = np.array( - [t.mask for t in tracks], dtype=bool - ) + detections.mask = np.array([t.mask for t in tracks], dtype=bool) else: detections.tracker_id = np.array([], dtype=int) return detections - def update_with_tensors(self, tensors: np.ndarray, masks:np.ndarray = None) -> List[STrack]: + def update_with_tensors( + self, tensors: np.ndarray, masks: np.ndarray = None + ) -> List[STrack]: """ Updates the tracker with the provided tensors and returns the updated tracks. @@ -293,8 +292,8 @@ class ByteTrack: masks_hs = masks[remain_inds] masks_ls = masks[inds_second] else: - masks_hs = np.array([None]*len(remain_inds)) - masks_ls = np.array([None]*len(inds_second)) + masks_hs = np.array([None] * len(remain_inds)) + masks_ls = np.array([None] * len(inds_second)) scores_keep = scores[remain_inds] scores_second = scores[inds_second] @@ -304,8 +303,8 @@ class ByteTrack: if len(dets) > 0: """Detections""" detections = [ - STrack(STrack.tlbr_to_tlwh(tlbr), s, c,m) - for (tlbr, s, c,m) in zip(dets, scores_keep, class_ids_keep,masks_hs) + STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m) + for (tlbr, s, c, m) in zip(dets, scores_keep, class_ids_keep, masks_hs) ] else: detections = [] @@ -345,8 +344,10 @@ class ByteTrack: if len(dets_second) > 0: """Detections""" detections_second = [ - STrack(STrack.tlbr_to_tlwh(tlbr), s, c,m) - for (tlbr, s, c, m) in zip(dets_second, scores_second, class_ids_second,masks_ls) + STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m) + for (tlbr, s, c, m) in zip( + dets_second, scores_second, class_ids_second, masks_ls + ) ] else: detections_second = [] From 30994243aed9044331090c8a5e7cc22434fa9c2a Mon Sep 17 00:00:00 2001 From: Antonio Consiglio <77753494+AntonioConsiglio@users.noreply.github.com> Date: Sat, 7 Oct 2023 19:15:03 +0200 Subject: [PATCH 3/6] Make sv.ByteTrack work with segmentation model#418 --- supervision/tracker/byte_tracker/core.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index 8ddf0963..e48c8396 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -75,7 +75,7 @@ class STrack(BaseTrack): if new_id: self.track_id = self.next_id() self.score = new_track.score - self.mask = new_track.mask if new_track.mask is not None else None + self.mask = new_track.mask def update(self, new_track, frame_id): """ @@ -97,8 +97,7 @@ class STrack(BaseTrack): self.score = new_track.score - self.mask = new_track.mask if new_track.mask is not None else None - + self.mask = new_track.mask @property def tlwh(self): """Get current position in bounding box format `(top left x, top left y, From 5efd535e9d9901b446d5cb690348a2de8cae1ef5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 7 Oct 2023 17:20:30 +0000 Subject: [PATCH 4/6] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/tracker/byte_tracker/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index a6e15c0c..82c932d6 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -97,7 +97,8 @@ class STrack(BaseTrack): self.score = new_track.score - self.mask = new_track.mask + self.mask = new_track.mask + @property def tlwh(self): """Get current position in bounding box format `(top left x, top left y, From 8bca45bde4ec6727fab80f52c49d94c4d47e3d3c Mon Sep 17 00:00:00 2001 From: Antonio Consiglio <77753494+AntonioConsiglio@users.noreply.github.com> Date: Mon, 9 Oct 2023 13:13:28 +0200 Subject: [PATCH 5/6] style and doc change as requested --- supervision/tracker/byte_tracker/core.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index a6e15c0c..fee8cbb1 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -1,4 +1,4 @@ -from typing import List, Tuple +from typing import List, Tuple, Optional import numpy as np @@ -11,7 +11,7 @@ from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter class STrack(BaseTrack): shared_kalman = KalmanFilter() - def __init__(self, tlwh, score, class_ids, mask=None): + def __init__(self, tlwh, score, class_ids, mask:Optional[np.array] = None): # wait activate self._tlwh = np.asarray(tlwh, dtype=np.float32) self.kalman_filter = None @@ -259,13 +259,14 @@ class ByteTrack: return detections def update_with_tensors( - self, tensors: np.ndarray, masks: np.ndarray = None + self, tensors: np.ndarray, masks: Optional[np.array] = None ) -> List[STrack]: """ Updates the tracker with the provided tensors and returns the updated tracks. Parameters: tensors: The new tensors to update with. + masks: The new masks associated to new tensors Returns: List[STrack]: Updated tracks. @@ -288,11 +289,11 @@ class ByteTrack: dets_second = bboxes[inds_second] dets = bboxes[remain_inds] if masks is not None: - masks_hs = masks[remain_inds] - masks_ls = masks[inds_second] + masks_keep = masks[remain_inds] + masks_second = masks[inds_second] else: - masks_hs = np.array([None] * len(remain_inds)) - masks_ls = np.array([None] * len(inds_second)) + masks_keep = np.array([None] * len(remain_inds)) + masks_second = np.array([None] * len(inds_second)) scores_keep = scores[remain_inds] scores_second = scores[inds_second] @@ -303,7 +304,7 @@ class ByteTrack: """Detections""" detections = [ STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m) - for (tlbr, s, c, m) in zip(dets, scores_keep, class_ids_keep, masks_hs) + for (tlbr, s, c, m) in zip(dets, scores_keep, class_ids_keep, masks_keep) ] else: detections = [] @@ -345,7 +346,7 @@ class ByteTrack: detections_second = [ STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m) for (tlbr, s, c, m) in zip( - dets_second, scores_second, class_ids_second, masks_ls + dets_second, scores_second, class_ids_second, masks_second ) ] else: From 49d268c25018b1a876b76a3cc2c7d0976e4ff421 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 9 Oct 2023 11:13:47 +0000 Subject: [PATCH 6/6] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/tracker/byte_tracker/core.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index 56ef4f45..03589c3d 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -1,4 +1,4 @@ -from typing import List, Tuple, Optional +from typing import List, Optional, Tuple import numpy as np @@ -11,7 +11,7 @@ from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter class STrack(BaseTrack): shared_kalman = KalmanFilter() - def __init__(self, tlwh, score, class_ids, mask:Optional[np.array] = None): + def __init__(self, tlwh, score, class_ids, mask: Optional[np.array] = None): # wait activate self._tlwh = np.asarray(tlwh, dtype=np.float32) self.kalman_filter = None @@ -305,7 +305,9 @@ class ByteTrack: """Detections""" detections = [ STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m) - for (tlbr, s, c, m) in zip(dets, scores_keep, class_ids_keep, masks_keep) + for (tlbr, s, c, m) in zip( + dets, scores_keep, class_ids_keep, masks_keep + ) ] else: detections = []