From fd1277f29a6cf503991f254648bb9d3feb7d2978 Mon Sep 17 00:00:00 2001 From: LinasKo Date: Fri, 1 Nov 2024 20:33:46 +0200 Subject: [PATCH] Add types, move STrack into a new file --- supervision/tracker/byte_tracker/core.py | 206 +----------------- supervision/tracker/byte_tracker/matching.py | 6 +- .../byte_tracker/single_object_track.py | 177 +++++++++++++++ supervision/tracker/byte_tracker/utils.py | 18 ++ 4 files changed, 209 insertions(+), 198 deletions(-) create mode 100644 supervision/tracker/byte_tracker/single_object_track.py create mode 100644 supervision/tracker/byte_tracker/utils.py diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index 019ce890..ecd844ad 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -1,4 +1,3 @@ -from enum import Enum from typing import List, Tuple import numpy as np @@ -7,194 +6,8 @@ from supervision.detection.core import Detections from supervision.detection.utils import box_iou_batch from supervision.tracker.byte_tracker import matching from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter - - -class TrackState(Enum): - New = 0 - Tracked = 1 - Lost = 2 - Removed = 3 - - -class IdCounter: - def __init__(self, start_id: int = 0): - self.start_id = start_id - if self.start_id <= self.NO_ID: - raise ValueError("start_id must be greater than -1") - self.reset() - - def reset(self) -> None: - self._id = self.start_id - - def new_id(self) -> int: - returned_id = self._id - self._id += 1 - return returned_id - - @property - def NO_ID(self) -> int: - return -1 - - -class STrack: - def __init__( - self, - tlwh, - score, - minimum_consecutive_frames, - shared_kalman: KalmanFilter, - internal_id_counter: IdCounter, - external_id_counter: IdCounter, - ): - self.state = TrackState.New - self.is_activated = False - self.start_frame = 0 - self.frame_id = 0 - - self._tlwh = np.asarray(tlwh, dtype=np.float32) - self.kalman_filter = None - self.shared_kalman = shared_kalman - self.mean, self.covariance = None, None - self.is_activated = False - - self.score = score - self.tracklet_len = 0 - - self.minimum_consecutive_frames = minimum_consecutive_frames - - self.internal_id_counter = internal_id_counter - self.external_id_counter = external_id_counter - self.internal_track_id = self.internal_id_counter.NO_ID - self.external_track_id = self.external_id_counter.NO_ID - - def predict(self): - mean_state = self.mean.copy() - if self.state != TrackState.Tracked: - mean_state[7] = 0 - self.mean, self.covariance = self.kalman_filter.predict( - mean_state, self.covariance - ) - - @staticmethod - def multi_predict(stracks, shared_kalman: KalmanFilter): - if len(stracks) > 0: - multi_mean = [] - multi_covariance = [] - for i, st in enumerate(stracks): - multi_mean.append(st.mean.copy()) - multi_covariance.append(st.covariance) - if st.state != TrackState.Tracked: - multi_mean[i][7] = 0 - - multi_mean, multi_covariance = shared_kalman.multi_predict( - np.asarray(multi_mean), np.asarray(multi_covariance) - ) - for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)): - stracks[i].mean = mean - stracks[i].covariance = cov - - def activate(self, kalman_filter, frame_id): - """Start a new tracklet""" - self.kalman_filter = kalman_filter - self.internal_track_id = self.internal_id_counter.new_id() - self.mean, self.covariance = self.kalman_filter.initiate( - self.tlwh_to_xyah(self._tlwh) - ) - - self.tracklet_len = 0 - self.state = TrackState.Tracked - if frame_id == 1: - self.is_activated = True - - if self.minimum_consecutive_frames == 1: - self.external_track_id = self.external_id_counter.new_id() - - self.frame_id = frame_id - self.start_frame = frame_id - - def re_activate(self, new_track, frame_id): - self.mean, self.covariance = self.kalman_filter.update( - self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh) - ) - self.tracklet_len = 0 - self.state = TrackState.Tracked - - self.frame_id = frame_id - self.score = new_track.score - - def update(self, new_track, frame_id): - """ - Update a matched track - :type new_track: STrack - :type frame_id: int - :type update_feature: bool - :return: - """ - self.frame_id = frame_id - self.tracklet_len += 1 - - new_tlwh = new_track.tlwh - self.mean, self.covariance = self.kalman_filter.update( - self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh) - ) - self.state = TrackState.Tracked - if self.tracklet_len == self.minimum_consecutive_frames: - self.is_activated = True - if self.external_track_id == self.external_id_counter.NO_ID: - self.external_track_id = self.external_id_counter.new_id() - - self.score = new_track.score - - @property - def tlwh(self): - """Get current position in bounding box format `(top left x, top left y, - width, height)`. - """ - if self.mean is None: - return self._tlwh.copy() - ret = self.mean[:4].copy() - ret[2] *= ret[3] - ret[:2] -= ret[2:] / 2 - return ret - - @property - def tlbr(self): - """Convert bounding box to format `(min x, min y, max x, max y)`, i.e., - `(top left, bottom right)`. - """ - ret = self.tlwh.copy() - ret[2:] += ret[:2] - return ret - - @staticmethod - def tlwh_to_xyah(tlwh): - """Convert bounding box to format `(center x, center y, aspect ratio, - height)`, where the aspect ratio is `width / height`. - """ - ret = np.asarray(tlwh).copy() - ret[:2] += ret[2:] / 2 - ret[2] /= ret[3] - return ret - - def to_xyah(self): - return self.tlwh_to_xyah(self.tlwh) - - @staticmethod - def tlbr_to_tlwh(tlbr): - ret = np.asarray(tlbr).copy() - ret[2:] -= ret[:2] - return ret - - @staticmethod - def tlwh_to_tlbr(tlwh): - ret = np.asarray(tlwh).copy() - ret[2:] += ret[:2] - return ret - - def __repr__(self): - return "OT_{}_({}-{})".format( - self.internal_track_id, self.start_frame, self.frame_id - ) +from supervision.tracker.byte_tracker.single_object_track import STrack, TrackState +from supervision.tracker.byte_tracker.utils import IdCounter class ByteTrack: @@ -291,7 +104,6 @@ class ByteTrack: ) ``` """ - tensors = np.hstack( ( detections.xyxy, @@ -323,7 +135,7 @@ class ByteTrack: return detections - def reset(self): + def reset(self) -> None: """ Resets the internal state of the ByteTrack tracker. @@ -335,9 +147,9 @@ class ByteTrack: self.frame_id = 0 self.internal_id_counter.reset() self.external_id_counter.reset() - self.tracked_tracks: List[STrack] = [] - self.lost_tracks: List[STrack] = [] - self.removed_tracks: List[STrack] = [] + self.tracked_tracks = [] + self.lost_tracks = [] + self.removed_tracks = [] def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]: """ @@ -529,7 +341,7 @@ def joint_tracks( return result -def sub_tracks(track_list_a: List, track_list_b: List) -> List[int]: +def sub_tracks(track_list_a: List[STrack], track_list_b: List[STrack]) -> List[int]: """ Returns a list of tracks from track_list_a after removing any tracks that share the same internal_track_id with tracks in track_list_b. @@ -550,7 +362,9 @@ def sub_tracks(track_list_a: List, track_list_b: List) -> List[int]: return list(tracks.values()) -def remove_duplicate_tracks(tracks_a: List, tracks_b: List) -> Tuple[List, List]: +def remove_duplicate_tracks( + tracks_a: List[STrack], tracks_b: List[STrack] +) -> Tuple[List[STrack], List[STrack]]: pairwise_distance = matching.iou_distance(tracks_a, tracks_b) matching_pairs = np.where(pairwise_distance < 0.15) diff --git a/supervision/tracker/byte_tracker/matching.py b/supervision/tracker/byte_tracker/matching.py index 4a95bc2c..eb774d4c 100644 --- a/supervision/tracker/byte_tracker/matching.py +++ b/supervision/tracker/byte_tracker/matching.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import TYPE_CHECKING, List, Tuple import numpy as np @@ -38,7 +40,7 @@ def linear_assignment( return indices_to_matches(cost_matrix, indices, thresh) -def iou_distance(atracks: List, btracks: List) -> np.ndarray: +def iou_distance(atracks: List[STrack], btracks: List[STrack]) -> np.ndarray: if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or ( len(btracks) > 0 and isinstance(btracks[0], np.ndarray) ): @@ -56,7 +58,7 @@ def iou_distance(atracks: List, btracks: List) -> np.ndarray: return cost_matrix -def fuse_score(cost_matrix: np.ndarray, stracks: List["STrack"]) -> np.ndarray: +def fuse_score(cost_matrix: np.ndarray, stracks: List[STrack]) -> np.ndarray: if cost_matrix.size == 0: return cost_matrix iou_sim = 1 - cost_matrix diff --git a/supervision/tracker/byte_tracker/single_object_track.py b/supervision/tracker/byte_tracker/single_object_track.py new file mode 100644 index 00000000..685b86e3 --- /dev/null +++ b/supervision/tracker/byte_tracker/single_object_track.py @@ -0,0 +1,177 @@ +from __future__ import annotations +from enum import Enum +from typing import List + +import numpy as np +import numpy.typing as npt + +from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter +from supervision.tracker.byte_tracker.utils import IdCounter + + +class TrackState(Enum): + New = 0 + Tracked = 1 + Lost = 2 + Removed = 3 + +class STrack: + def __init__( + self, + tlwh: npt.NDArray[np.float32], + score: npt.NDArray[np.float32], + minimum_consecutive_frames: int, + shared_kalman: KalmanFilter, + internal_id_counter: IdCounter, + external_id_counter: IdCounter, + ): + self.state = TrackState.New + self.is_activated = False + self.start_frame = 0 + self.frame_id = 0 + + self._tlwh = np.asarray(tlwh, dtype=np.float32) + self.kalman_filter = None + self.shared_kalman = shared_kalman + self.mean, self.covariance = None, None + self.is_activated = False + + self.score = score + self.tracklet_len = 0 + + self.minimum_consecutive_frames = minimum_consecutive_frames + + self.internal_id_counter = internal_id_counter + self.external_id_counter = external_id_counter + self.internal_track_id = self.internal_id_counter.NO_ID + self.external_track_id = self.external_id_counter.NO_ID + + def predict(self) -> None: + mean_state = self.mean.copy() + if self.state != TrackState.Tracked: + mean_state[7] = 0 + self.mean, self.covariance = self.kalman_filter.predict( + mean_state, self.covariance + ) + + @staticmethod + def multi_predict(stracks: List[STrack], shared_kalman: KalmanFilter) -> None: + if len(stracks) > 0: + multi_mean = [] + multi_covariance = [] + for i, st in enumerate(stracks): + multi_mean.append(st.mean.copy()) + multi_covariance.append(st.covariance) + if st.state != TrackState.Tracked: + multi_mean[i][7] = 0 + + multi_mean, multi_covariance = shared_kalman.multi_predict( + np.asarray(multi_mean), np.asarray(multi_covariance) + ) + for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)): + stracks[i].mean = mean + stracks[i].covariance = cov + + def activate(self, kalman_filter: KalmanFilter, frame_id: int) -> None: + """Start a new tracklet""" + self.kalman_filter = kalman_filter + self.internal_track_id = self.internal_id_counter.new_id() + self.mean, self.covariance = self.kalman_filter.initiate( + self.tlwh_to_xyah(self._tlwh) + ) + + self.tracklet_len = 0 + self.state = TrackState.Tracked + if frame_id == 1: + self.is_activated = True + + if self.minimum_consecutive_frames == 1: + self.external_track_id = self.external_id_counter.new_id() + + self.frame_id = frame_id + self.start_frame = frame_id + + def re_activate(self, new_track: STrack, frame_id: int) -> None: + self.mean, self.covariance = self.kalman_filter.update( + self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh) + ) + self.tracklet_len = 0 + self.state = TrackState.Tracked + + self.frame_id = frame_id + self.score = new_track.score + + def update(self, new_track: STrack, frame_id: int) -> None: + """ + Update a matched track + :type new_track: STrack + :type frame_id: int + :type update_feature: bool + :return: + """ + self.frame_id = frame_id + self.tracklet_len += 1 + + new_tlwh = new_track.tlwh + self.mean, self.covariance = self.kalman_filter.update( + self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh) + ) + self.state = TrackState.Tracked + if self.tracklet_len == self.minimum_consecutive_frames: + self.is_activated = True + if self.external_track_id == self.external_id_counter.NO_ID: + self.external_track_id = self.external_id_counter.new_id() + + self.score = new_track.score + + @property + def tlwh(self) -> npt.NDArray[np.float32]: + """Get current position in bounding box format `(top left x, top left y, + width, height)`. + """ + if self.mean is None: + return self._tlwh.copy() + ret = self.mean[:4].copy() + ret[2] *= ret[3] + ret[:2] -= ret[2:] / 2 + return ret + + @property + def tlbr(self) -> npt.NDArray[np.float32]: + """Convert bounding box to format `(min x, min y, max x, max y)`, i.e., + `(top left, bottom right)`. + """ + ret = self.tlwh.copy() + ret[2:] += ret[:2] + return ret + + @staticmethod + def tlwh_to_xyah(tlwh) -> npt.NDArray[np.float32]: + """Convert bounding box to format `(center x, center y, aspect ratio, + height)`, where the aspect ratio is `width / height`. + """ + ret = np.asarray(tlwh).copy() + ret[:2] += ret[2:] / 2 + ret[2] /= ret[3] + return ret + + def to_xyah(self) -> npt.NDArray[np.float32]: + return self.tlwh_to_xyah(self.tlwh) + + @staticmethod + def tlbr_to_tlwh(tlbr) -> npt.NDArray[np.float32]: + ret = np.asarray(tlbr).copy() + ret[2:] -= ret[:2] + return ret + + @staticmethod + def tlwh_to_tlbr(tlwh) -> npt.NDArray[np.float32]: + ret = np.asarray(tlwh).copy() + ret[2:] += ret[:2] + return ret + + def __repr__(self) -> str: + return "OT_{}_({}-{})".format( + self.internal_track_id, self.start_frame, self.frame_id + ) + diff --git a/supervision/tracker/byte_tracker/utils.py b/supervision/tracker/byte_tracker/utils.py new file mode 100644 index 00000000..b7a62645 --- /dev/null +++ b/supervision/tracker/byte_tracker/utils.py @@ -0,0 +1,18 @@ +class IdCounter: + def __init__(self, start_id: int = 0): + self.start_id = start_id + if self.start_id <= self.NO_ID: + raise ValueError("start_id must be greater than -1") + self.reset() + + def reset(self) -> None: + self._id = self.start_id + + def new_id(self) -> int: + returned_id = self._id + self._id += 1 + return returned_id + + @property + def NO_ID(self) -> int: + return -1 \ No newline at end of file