🧹 more cleanup - remove more unused code, documentation improvements
This commit is contained in:
SkalskiP
parent
ba880402ad
commit
db2eede55b
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
... )
|
||||
```
|
||||
"""
|
||||
|
|
|
|||
Loading…
Reference in New Issue