From db2eede55b98f6a732b4a1cdda47d6cbcb7c331b Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 7 Aug 2023 14:39:25 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B9=20more=20cleanup=20-=20remove=20mo?= =?UTF-8?q?re=20unused=20code,=20documentation=20improvements?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/tracker/byte_tracker/core.py | 39 ++++-- supervision/tracker/byte_tracker/matching.py | 130 +------------------ supervision/utils/video.py | 10 +- 3 files changed, 37 insertions(+), 142 deletions(-) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index 6a23e918..4ca6a489 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -2,7 +2,7 @@ from typing import List, Tuple import numpy as np -from supervision import Detections +from supervision.detection.core import Detections from supervision.tracker.byte_tracker import matching from supervision.tracker.byte_tracker.basetrack import BaseTrack, TrackState from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter @@ -143,16 +143,13 @@ class STrack(BaseTrack): return "OT_{}_({}-{})".format(self.track_id, self.start_frame, self.end_frame) -# converts Detections into format that can be consumed by match_detections_with_tracks function def detections2boxes(detections: Detections) -> np.ndarray: """ - Convert Detections into a format that can be consumed by the match_detections_with_tracks function. - - Parameters: - detections (Detections): An object representing the detected bounding boxes. - + Convert Supervision Detections to numpy tensors for further computation. + Args: + detections (Detections): Detections/Targets in the format of sv.Detections. Returns: - np.ndarray: An array containing the bounding boxes' coordinates (xyxy) and their corresponding confidences. + (np.ndarray): Detections as numpy tensors as in `(x_min, y_min, x_max, y_max, confidence, class_id)` order. """ return np.hstack( ( @@ -201,8 +198,31 @@ class ByteTrack: detections: The new detections to update with. Returns: Detection: supervision detection result with track id. - Examples: + Example: ```python + >>> import supervision as sv + >>> from ultralytics import YOLO + + >>> model = YOLO(...) + >>> byte_tracker = sv.ByteTrack() + >>> annotator = sv.BoxAnnotator() + + def callback(frame: np.ndarray, index: int) -> np.ndarray: + results = model(frame)[0] + detections = sv.Detections.from_yolov8(results) + detections = byte_tracker.update_from_detections(detections=detections) + labels = [ + f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}" + for _, _, confidence, class_id, tracker_id + in detections + ] + return annotator.annotate(scene=frame.copy(), detections=detections, labels=labels) + + sv.process_video( + source_path='...', + target_path='...', + callback=callback + ) ``` """ @@ -225,7 +245,6 @@ class ByteTrack: Parameters: output_results: The new detections to update with. - Updates the strack with the provided results and frame info. Returns: Track_id: track id diff --git a/supervision/tracker/byte_tracker/matching.py b/supervision/tracker/byte_tracker/matching.py index 83ee5aa1..5a880bd5 100644 --- a/supervision/tracker/byte_tracker/matching.py +++ b/supervision/tracker/byte_tracker/matching.py @@ -1,32 +1,12 @@ -from typing import List, Optional, Tuple +from typing import List, Tuple import numpy as np -import scipy from scipy.optimize import linear_sum_assignment -from scipy.spatial.distance import cdist from supervision.detection.utils import box_iou_batch -from supervision.tracker.byte_tracker import kalman_filter -def merge_matches(m1, m2, shape) -> Tuple[List, tuple, tuple]: - O, P, Q = shape - m1 = np.asarray(m1) - m2 = np.asarray(m2) - - M1 = scipy.sparse.coo_matrix((np.ones(len(m1)), (m1[:, 0], m1[:, 1])), shape=(O, P)) - M2 = scipy.sparse.coo_matrix((np.ones(len(m2)), (m2[:, 0], m2[:, 1])), shape=(P, Q)) - - mask = M1 * M2 - match = mask.nonzero() - match = list(zip(match[0], match[1])) - unmatched_O = tuple(set(range(O)) - set([i for i, j in match])) - unmatched_Q = tuple(set(range(Q)) - set([j for i, j in match])) - - return match, unmatched_O, unmatched_Q - - -def _indices_to_matches( +def indices_to_matches( cost_matrix: np.ndarray, indices: np.ndarray, thresh: float ) -> Tuple[np.ndarray, tuple, tuple]: matched_cost = cost_matrix[tuple(zip(*indices))] @@ -58,7 +38,7 @@ def linear_assignment( row_ind, col_ind = linear_sum_assignment(cost_matrix) indices = np.column_stack((row_ind, col_ind)) - return _indices_to_matches(cost_matrix, indices, thresh) + return indices_to_matches(cost_matrix, indices, thresh) def iou_distance(atracks: List, btracks: List) -> np.ndarray: @@ -87,110 +67,6 @@ def iou_distance(atracks: List, btracks: List) -> np.ndarray: return cost_matrix -def v_iou_distance(atracks: List, btracks: List) -> np.ndarray: - """ - Compute cost based on IoU - :type atracks: list[STrack] - :type btracks: list[STrack] - - :rtype cost_matrix np.ndarray - """ - - if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or ( - len(btracks) > 0 and isinstance(btracks[0], np.ndarray) - ): - atlbrs = atracks - btlbrs = btracks - else: - atlbrs = [track.tlwh_to_tlbr(track.pred_bbox) for track in atracks] - btlbrs = [track.tlwh_to_tlbr(track.pred_bbox) for track in btracks] - _ious = box_iou_batch(np.asarray(atlbrs), np.asarray(btlbrs)) - cost_matrix = 1 - _ious - - return cost_matrix - - -def embedding_distance(tracks: List, detections: List, metric="cosine") -> np.ndarray: - """ - :param tracks: list[STrack] - :param detections: list[BaseTrack] - :param metric: - :return: cost_matrix np.ndarray - """ - - cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float32) - if cost_matrix.size == 0: - return cost_matrix - det_features = np.asarray( - [track.curr_feat for track in detections], dtype=np.float32 - ) - track_features = np.asarray( - [track.smooth_feat for track in tracks], dtype=np.float32 - ) - cost_matrix = np.maximum( - 0.0, cdist(track_features, det_features, metric) - ) # Nomalized features - return cost_matrix - - -def gate_cost_matrix( - kf, - cost_matrix: np.ndarray, - tracks: List, - detections: np.ndarray, - only_position=False, -) -> np.ndarray: - if cost_matrix.size == 0: - return cost_matrix - gating_dim = 2 if only_position else 4 - gating_threshold = kalman_filter.chi2inv95[gating_dim] - measurements = np.asarray([det.to_xyah() for det in detections]) - for row, track in enumerate(tracks): - gating_distance = kf.gating_distance( - track.mean, track.covariance, measurements, only_position - ) - cost_matrix[row, gating_distance > gating_threshold] = np.inf - return cost_matrix - - -def fuse_motion( - kf, - cost_matrix: np.ndarray, - tracks: List, - detections: np.ndarray, - only_position=False, - lambda_=0.98, -) -> np.ndarray: - if cost_matrix.size == 0: - return cost_matrix - gating_dim = 2 if only_position else 4 - gating_threshold = kalman_filter.chi2inv95[gating_dim] - measurements = np.asarray([det.to_xyah() for det in detections]) - for row, track in enumerate(tracks): - gating_distance = kf.gating_distance( - track.mean, track.covariance, measurements, only_position, metric="maha" - ) - cost_matrix[row, gating_distance > gating_threshold] = np.inf - cost_matrix[row] = lambda_ * cost_matrix[row] + (1 - lambda_) * gating_distance - return cost_matrix - - -def fuse_iou( - cost_matrix: np.ndarray, tracks: List, detections: np.ndarray -) -> np.ndarray: - if cost_matrix.size == 0: - return cost_matrix - reid_sim = 1 - cost_matrix - iou_dist = iou_distance(tracks, detections) - iou_sim = 1 - iou_dist - fuse_sim = reid_sim * (1 + iou_sim) / 2 - det_scores = np.array([det.score for det in detections]) - det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0) - # fuse_sim = fuse_sim * (1 + det_scores) / 2 - fuse_cost = 1 - fuse_sim - return fuse_cost - - def fuse_score(cost_matrix: np.ndarray, detections: List) -> np.ndarray: if cost_matrix.size == 0: return cost_matrix diff --git a/supervision/utils/video.py b/supervision/utils/video.py index fcec53b1..dda68ca6 100644 --- a/supervision/utils/video.py +++ b/supervision/utils/video.py @@ -163,15 +163,15 @@ def process_video( Examples: ```python - >>> from supervision import process_video + >>> import supervision as sv - >>> def process_frame(scene: np.ndarray) -> np.ndarray: + >>> def callback(scene: np.ndarray, index: int) -> np.ndarray: ... ... >>> process_video( - ... source_path='source_video.mp4', - ... target_path='target_video.mp4', - ... callback=process_frame + ... source_path='...', + ... target_path='...', + ... callback=callback ... ) ``` """