xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

updated tracker

This commit is contained in:
dbroboflow 2023-08-01 12:30:52 +00:00
parent 27e8e2d91b
commit a59e4dfe73
13 changed files with 1253 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
docs/tracker/bytetrack.md Normal file
View File

@ -0,0 +1,7 @@
## update_from_detections
:::supervision.tracker.byte_tracker.byte_tracker.ByteTrack.update_from_detections
## update_from_numpy
:::supervision.tracker.byte_tracker.byte_tracker.ByteTrack.update_from_numpy

BIN
docs/tracker/output.mp4 Normal file
View File

Binary file not shown.

57
docs/tracker/tracker_bytetrack_detections.py Normal file
View File

@ -0,0 +1,57 @@
from tqdm.notebook import tqdm
from ultralytics import YOLO
import supervision as sv
from supervision.detection.annotate import BoxAnnotator
from supervision import ByteTrack
from supervision.utils.video import VideoInfo, VideoSink, get_video_frames_generator
model = YOLO("yolov5s.pt")
CLASS_NAMES_DICT = model.model.names
SOURCE_VIDEO_PATH = "walking.mp4"
TARGET_VIDEO_PATH = "output.mp4"
byte_tracker = ByteTrack(
track_thresh=0.25,
track_buffer=30,
match_thresh=0.8,
aspect_ratio_thresh=3.0,
min_box_area=1.0,
mot20=False,
frame_rate=30,
)
annotator = BoxAnnotator()
# create VideoInfo instance
video_info = VideoInfo.from_video_path(SOURCE_VIDEO_PATH)
# create frame generator
generator = get_video_frames_generator(SOURCE_VIDEO_PATH)
# open target video file
with VideoSink(TARGET_VIDEO_PATH, video_info) as sink:
# loop over video frames
for frame in tqdm(generator, total=video_info.total_frames):
# inference
results = model(frame)
sv_results = sv.Detections.from_yolov8(results[0])
# update tracker
detections_res = byte_tracker.update_from_detections(
detections=sv_results, img_info=frame.shape, img_size=frame.shape
)
frame_text_list = []
if len(detections_res) > 0:
for res in detections_res:
confidence = res[2]
class_id = res[3]
tracker_id = res[4]
frame_text = (
f"#{tracker_id} {CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
)
frame_text_list.append(frame_text)
# draw bbox and add class name with track id
frame = annotator.annotate(
scene=frame, detections=detections_res, labels=frame_text_list
)
sink.write_frame(frame)

67
docs/tracker/tracker_bytetrack_numpy.py Normal file
View File

@ -0,0 +1,67 @@
import numpy as np
from tqdm.notebook import tqdm
from ultralytics import YOLO
import supervision as sv
from supervision.detection.annotate import BoxAnnotator
from supervision import (
ByteTrack,
detections2boxes,
match_detections_with_tracks,
)
from supervision.utils.video import VideoInfo, VideoSink, get_video_frames_generator
model = YOLO("yolov5s.pt")
CLASS_NAMES_DICT = model.model.names
SOURCE_VIDEO_PATH = "walking.mp4"
TARGET_VIDEO_PATH = "output.mp4"
byte_tracker = ByteTrack(
track_thresh=0.25,
track_buffer=30,
match_thresh=0.8,
aspect_ratio_thresh=3.0,
min_box_area=1.0,
mot20=False,
frame_rate=30,
)
annotator = BoxAnnotator()
# create VideoInfo instance
video_info = VideoInfo.from_video_path(SOURCE_VIDEO_PATH)
# create frame generator
generator = get_video_frames_generator(SOURCE_VIDEO_PATH)
# open target video file
with VideoSink(TARGET_VIDEO_PATH, video_info) as sink:
# loop over video frames
for frame in tqdm(generator, total=video_info.total_frames):
# inference
results = model(frame)
detections = sv.Detections.from_yolov8(results[0])
# update tracker
tracks = byte_tracker.update_from_numpy(
output_results=detections2boxes(detections=detections),
img_info=frame.shape,
img_size=frame.shape,
)
tracker_id = match_detections_with_tracks(detections=detections, tracks=tracks)
if tracker_id != None:
detections.tracker_id = np.array(tracker_id)
frame_text_list = []
if len(detections) > 0:
for res in detections:
confidence = res[2]
class_id = res[3]
tracker_id = res[4]
frame_text = (
f"#{tracker_id} {CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
)
frame_text_list.append(frame_text)
# draw bbox and add class name with track id
frame = annotator.annotate(
scene=frame, detections=detections, labels=frame_text_list
)
sink.write_frame(frame)

BIN
docs/tracker/walking.mp4 Normal file
View File

Binary file not shown.

BIN
docs/tracker/yolov5su.pt Normal file
View File

Binary file not shown.

5
supervision/__init__.py
View File

@ -40,3 +40,8 @@ from supervision.utils.video import (
get_video_frames_generator,
process_video,
)
from supervision.tracker.byte_tracker.byte_tracker import (
ByteTrack,
detections2boxes,
match_detections_with_tracks,
)

0
supervision/tracker/__init__.py Normal file
View File

0
supervision/tracker/byte_tracker/__init__.py Normal file
View File

55
supervision/tracker/byte_tracker/basetrack.py Normal file
View File

@ -0,0 +1,55 @@
from collections import OrderedDict
from enum import Enum
import numpy as np
class TrackState(Enum):
New = 0
Tracked = 1
Lost = 2
Removed = 3
class BaseTrack:
_count = 0
def __init__(self):
self.track_id = 0
self.is_activated = False
self.state = TrackState.New
self.history = OrderedDict()
self.features = []
self.curr_feature = None
self.score = 0
self.start_frame = 0
self.frame_id = 0
self.time_since_update = 0
# multi-camera
self.location = (np.inf, np.inf)
@property
def end_frame(self):
return self.frame_id
@staticmethod
def next_id():
BaseTrack._count += 1
return BaseTrack._count
def activate(self, *args):
raise NotImplementedError
def predict(self):
raise NotImplementedError
def update(self, *args, **kwargs):
raise NotImplementedError
def mark_lost(self):
self.state = TrackState.Lost
def mark_removed(self):
self.state = TrackState.Removed

541
supervision/tracker/byte_tracker/byte_tracker.py Normal file
View File

@ -0,0 +1,541 @@
from typing import List
import numpy as np
from supervision import Detections
from supervision.detection.utils import box_iou_batch
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
class Strack(BaseTrack):
shared_kalman = KalmanFilter()
def __init__(self, tlwh, score):
# wait activate
self._tlwh = np.asarray(tlwh, dtype=np.float32)
self.kalman_filter = None
self.mean, self.covariance = None, None
self.is_activated = False
self.score = score
self.tracklet_len = 0
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):
if len(stracks) > 0:
multi_mean = np.asarray([st.mean.copy() for st in stracks])
multi_covariance = np.asarray([st.covariance for st in stracks])
for i, st in enumerate(stracks):
if st.state != TrackState.Tracked:
multi_mean[i][7] = 0
multi_mean, multi_covariance = Strack.shared_kalman.multi_predict(
multi_mean, 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.track_id = self.next_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
# self.is_activated = True
self.frame_id = frame_id
self.start_frame = frame_id
def re_activate(self, new_track, frame_id, new_id=False):
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.is_activated = True
self.frame_id = frame_id
if new_id:
self.track_id = self.next_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
self.is_activated = True
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.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.
Returns:
np.ndarray: An array containing the bounding boxes' coordinates (xyxy) and their corresponding confidences.
"""
return np.hstack((detections.xyxy, detections.confidence[:, np.newaxis]))
# converts List[strack] into format that can be consumed by match_detections_with_tracks function
def tracks2boxes(tracks: List[Strack]) -> np.ndarray:
return np.array([track.tlbr for track in tracks], dtype=float)
# matches our bounding boxes with predictions
def match_detections_with_tracks(
detections: Detections, tracks: List[Strack]
) -> Detections:
"""
Match bounding boxes from detections with existing strack objects.
Parameters:
detections (Detections): An object representing the detected bounding boxes.
tracks (List[strack]): A list of strack objects to match with the detections.
Returns:
Detections: An object representing the matched tracker IDs for the detections.
"""
if not np.any(detections.xyxy) or len(tracks) == 0:
return np.empty((0,))
tracks_boxes = tracks2boxes(tracks=tracks)
iou = box_iou_batch(tracks_boxes, detections.xyxy)
track2detection = np.argmax(iou, axis=1)
tracker_ids = [None] * len(detections)
for tracker_index, detection_index in enumerate(track2detection):
if iou[tracker_index, detection_index] != 0:
tracker_ids[detection_index] = tracks[tracker_index].track_id
return tracker_ids
class ByteTrack:
def __init__(
self,
track_thresh=0.25,
track_buffer=30,
match_thresh=0.8,
aspect_ratio_thresh=3.0,
min_box_area=1.0,
mot20=False,
frame_rate=30,
):
"""
Initialize the ByteTrack object.
Parameters:
track_thresh (float, optional): Detection confidence threshold for track activation.
track_buffer (int, optional): Number of frames to buffer when a track is lost.
match_thresh (float, optional): Threshold for matching tracks with detections.
aspect_ratio_thresh (float, optional): Threshold for aspect ratio to filter out detections.
min_box_area (float, optional): Minimum box area threshold to filter out detections.
mot20 (bool, optional): Set to True for using MOT20 evaluation criteria.
frame_rate (int, optional): The frame rate of the video.
"""
self.track_thresh = track_thresh
self.track_buffer = track_buffer
self.match_thresh = match_thresh
self.aspect_ratio_thresh = aspect_ratio_thresh
self.min_box_area = min_box_area
self.mot20 = mot20
self.tracked_stracks = [] # type: list[Strack]
self.lost_stracks = [] # type: list[Strack]
self.removed_stracks = [] # type: list[Strack]
self.frame_id = 0
self.det_thresh = self.track_thresh + 0.1
self.buffer_size = int(frame_rate / 30.0 * self.track_buffer)
self.max_time_lost = self.buffer_size
self.kalman_filter = KalmanFilter()
def update_from_detections(self, detections, img_info, img_size):
"""
Updates the strack with the provided results and frame info.
Parameters:
detections: The new detections to update with.
img_info: Image Info.
img_size: Image Size
Returns:
Detection: supervision detection result with track id
Examples:
```python
from tqdm.notebook import tqdm
from ultralytics import YOLO
import supervision as sv
from supervision.detection.annotate import BoxAnnotator
from supervision import ByteTrack
from supervision.utils.video import VideoInfo, VideoSink, get_video_frames_generator
model = YOLO("yolov5s.pt")
CLASS_NAMES_DICT = model.model.names
SOURCE_VIDEO_PATH = "walking.mp4"
TARGET_VIDEO_PATH = "output.mp4"
byte_tracker = ByteTrack(
track_thresh=0.25,
track_buffer=30,
match_thresh=0.8,
aspect_ratio_thresh=3.0,
min_box_area=1.0,
mot20=False,
frame_rate=30,
)
annotator = BoxAnnotator()
# create VideoInfo instance
video_info = VideoInfo.from_video_path(SOURCE_VIDEO_PATH)
# create frame generator
generator = get_video_frames_generator(SOURCE_VIDEO_PATH)
# open target video file
with VideoSink(TARGET_VIDEO_PATH, video_info) as sink:
# loop over video frames
for frame in tqdm(generator, total=video_info.total_frames):
# inference
results = model(frame)
sv_results = sv.Detections.from_yolov8(results[0])
# update tracker
detections_res = byte_tracker.update_from_detections(
detections=sv_results, img_info=frame.shape, img_size=frame.shape
)
frame_text_list = []
if len(detections_res) > 0:
for res in detections_res:
confidence = res[2]
class_id = res[3]
tracker_id = res[4]
frame_text = (
f"#{tracker_id} {CLASS_NAMES_DICT[class_id]} {confidence:0.2f}"
)
frame_text_list.append(frame_text)
# draw bbox and add class name with track id
frame = annotator.annotate(
scene=frame, detections=detections_res, labels=frame_text_list
)
sink.write_frame(frame)
```
"""
# update tracker
tracks = self.update_from_numpy(
output_results=detections2boxes(detections=detections),
img_info=img_info,
img_size=img_size,
)
tracker_id = match_detections_with_tracks(detections=detections, tracks=tracks)
if tracker_id != None:
detections.tracker_id = np.array(tracker_id)
return detections
def update_from_numpy(self, output_results, img_info, img_size):
"""
Update a matched strack. It uses the numpy results.
Parameters:
output_results: The new detections to update with.
img_info: Image Info.
img_size: Image Size
Updates the strack with the provided results and frame info.
Returns:
Track_id: track id
Examples:
```python
from supervision import (
ByteTrack,
detections2boxes,
match_detections_with_tracks,
)
# update tracker
tracks = byte_tracker.update_from_numpy(
output_results=detections2boxes(detections=detections),
img_info=frame.shape,
img_size=frame.shape,
)
tracker_id = match_detections_with_tracks(detections=detections, tracks=tracks)
```
"""
self.frame_id += 1
activated_starcks = []
refind_stracks = []
lost_stracks = []
removed_stracks = []
if output_results.shape[1] == 5:
scores = output_results[:, 4]
bboxes = output_results[:, :4]
else:
output_results = output_results.cpu().numpy()
scores = output_results[:, 4] * output_results[:, 5]
bboxes = output_results[:, :4] # x1y1x2y2
img_h, img_w = img_info[0], img_info[1]
scale = min(img_size[0] / float(img_h), img_size[1] / float(img_w))
bboxes /= scale
remain_inds = scores > self.track_thresh
inds_low = scores > 0.1
inds_high = scores < self.track_thresh
inds_second = np.logical_and(inds_low, inds_high)
dets_second = bboxes[inds_second]
dets = bboxes[remain_inds]
scores_keep = scores[remain_inds]
scores_second = scores[inds_second]
if len(dets) > 0:
"""Detections"""
detections = [
Strack(Strack.tlbr_to_tlwh(tlbr), s)
for (tlbr, s) in zip(dets, scores_keep)
]
else:
detections = []
""" Add newly detected tracklets to tracked_stracks"""
unconfirmed = []
tracked_stracks = [] # type: list[Strack]
for track in self.tracked_stracks:
if not track.is_activated:
unconfirmed.append(track)
else:
tracked_stracks.append(track)
""" Step 2: First association, with high score detection boxes"""
strack_pool = joint_stracks(tracked_stracks, self.lost_stracks)
# Predict the current location with KF
Strack.multi_predict(strack_pool)
dists = matching.iou_distance(strack_pool, detections)
if not self.mot20:
dists = matching.fuse_score(dists, detections)
matches, u_track, u_detection = matching.linear_assignment(
dists, thresh=self.match_thresh
)
for itracked, idet in matches:
track = strack_pool[itracked]
det = detections[idet]
if track.state == TrackState.Tracked:
track.update(detections[idet], self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
""" Step 3: Second association, with low score detection boxes"""
# association the untrack to the low score detections
if len(dets_second) > 0:
"""Detections"""
detections_second = [
Strack(Strack.tlbr_to_tlwh(tlbr), s)
for (tlbr, s) in zip(dets_second, scores_second)
]
else:
detections_second = []
r_tracked_stracks = [
strack_pool[i]
for i in u_track
if strack_pool[i].state == TrackState.Tracked
]
dists = matching.iou_distance(r_tracked_stracks, detections_second)
matches, u_track, u_detection_second = matching.linear_assignment(
dists, thresh=0.5
)
for itracked, idet in matches:
track = r_tracked_stracks[itracked]
det = detections_second[idet]
if track.state == TrackState.Tracked:
track.update(det, self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
for it in u_track:
track = r_tracked_stracks[it]
if not track.state == TrackState.Lost:
track.mark_lost()
lost_stracks.append(track)
"""Deal with unconfirmed tracks, usually tracks with only one beginning frame"""
detections = [detections[i] for i in u_detection]
dists = matching.iou_distance(unconfirmed, detections)
if not self.mot20:
dists = matching.fuse_score(dists, detections)
matches, u_unconfirmed, u_detection = matching.linear_assignment(
dists, thresh=0.7
)
for itracked, idet in matches:
unconfirmed[itracked].update(detections[idet], self.frame_id)
activated_starcks.append(unconfirmed[itracked])
for it in u_unconfirmed:
track = unconfirmed[it]
track.mark_removed()
removed_stracks.append(track)
""" Step 4: Init new stracks"""
for inew in u_detection:
track = detections[inew]
if track.score < self.det_thresh:
continue
track.activate(self.kalman_filter, self.frame_id)
activated_starcks.append(track)
""" Step 5: Update state"""
for track in self.lost_stracks:
if self.frame_id - track.end_frame > self.max_time_lost:
track.mark_removed()
removed_stracks.append(track)
# print('Ramained match {} s'.format(t4-t3))
self.tracked_stracks = [
t for t in self.tracked_stracks if t.state == TrackState.Tracked
]
self.tracked_stracks = joint_stracks(self.tracked_stracks, activated_starcks)
self.tracked_stracks = joint_stracks(self.tracked_stracks, refind_stracks)
self.lost_stracks = sub_stracks(self.lost_stracks, self.tracked_stracks)
self.lost_stracks.extend(lost_stracks)
self.lost_stracks = sub_stracks(self.lost_stracks, self.removed_stracks)
self.removed_stracks.extend(removed_stracks)
self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(
self.tracked_stracks, self.lost_stracks
)
# get scores of lost tracks
output_stracks = [track for track in self.tracked_stracks if track.is_activated]
return output_stracks
def joint_stracks(tlista, tlistb):
exists = {}
res = []
for t in tlista:
exists[t.track_id] = 1
res.append(t)
for t in tlistb:
tid = t.track_id
if not exists.get(tid, 0):
exists[tid] = 1
res.append(t)
return res
def sub_stracks(tlista, tlistb):
stracks = {}
for t in tlista:
stracks[t.track_id] = t
for t in tlistb:
tid = t.track_id
if stracks.get(tid, 0):
del stracks[tid]
return list(stracks.values())
def remove_duplicate_stracks(stracksa, stracksb):
pdist = matching.iou_distance(stracksa, stracksb)
pairs = np.where(pdist < 0.15)
dupa, dupb = list(), list()
for p, q in zip(*pairs):
timep = stracksa[p].frame_id - stracksa[p].start_frame
timeq = stracksb[q].frame_id - stracksb[q].start_frame
if timep > timeq:
dupb.append(q)
else:
dupa.append(p)
resa = [t for i, t in enumerate(stracksa) if not i in dupa]
resb = [t for i, t in enumerate(stracksb) if not i in dupb]
return resa, resb

284
supervision/tracker/byte_tracker/kalman_filter.py Normal file
View File

@ -0,0 +1,284 @@
# vim: expandtab:ts=4:sw=4
import numpy as np
import scipy.linalg
"""
Table for the 0.95 quantile of the chi-square distribution with N degrees of
freedom (contains values for N=1, ..., 9). Taken from MATLAB/Octave's chi2inv
function and used as Mahalanobis gating threshold.
"""
chi2inv95 = {
1: 3.8415,
2: 5.9915,
3: 7.8147,
4: 9.4877,
5: 11.070,
6: 12.592,
7: 14.067,
8: 15.507,
9: 16.919,
}
class KalmanFilter(object):
"""
A simple Kalman filter for tracking bounding boxes in image space.
The 8-dimensional state space
x, y, a, h, vx, vy, va, vh
contains the bounding box center position (x, y), aspect ratio a, height h,
and their respective velocities.
Object motion follows a constant velocity model. The bounding box location
(x, y, a, h) is taken as direct observation of the state space (linear
observation model).
"""
def __init__(self):
ndim, dt = 4, 1.0
# Create Kalman filter model matrices.
self._motion_mat = np.eye(2 * ndim, 2 * ndim)
for i in range(ndim):
self._motion_mat[i, ndim + i] = dt
self._update_mat = np.eye(ndim, 2 * ndim)
# Motion and observation uncertainty are chosen relative to the current
# state estimate. These weights control the amount of uncertainty in
# the model. This is a bit hacky.
self._std_weight_position = 1.0 / 20
self._std_weight_velocity = 1.0 / 160
def initiate(self, measurement):
"""Create track from unassociated measurement.
Parameters
----------
measurement : ndarray
Bounding box coordinates (x, y, a, h) with center position (x, y),
aspect ratio a, and height h.
Returns
-------
(ndarray, ndarray)
Returns the mean vector (8 dimensional) and covariance matrix (8x8
dimensional) of the new track. Unobserved velocities are initialized
to 0 mean.
"""
mean_pos = measurement
mean_vel = np.zeros_like(mean_pos)
mean = np.r_[mean_pos, mean_vel]
std = [
2 * self._std_weight_position * measurement[3],
2 * self._std_weight_position * measurement[3],
1e-2,
2 * self._std_weight_position * measurement[3],
10 * self._std_weight_velocity * measurement[3],
10 * self._std_weight_velocity * measurement[3],
1e-5,
10 * self._std_weight_velocity * measurement[3],
]
covariance = np.diag(np.square(std))
return mean, covariance
def predict(self, mean, covariance):
"""Run Kalman filter prediction step.
Parameters
----------
mean : ndarray
The 8 dimensional mean vector of the object state at the previous
time step.
covariance : ndarray
The 8x8 dimensional covariance matrix of the object state at the
previous time step.
Returns
-------
(ndarray, ndarray)
Returns the mean vector and covariance matrix of the predicted
state. Unobserved velocities are initialized to 0 mean.
"""
std_pos = [
self._std_weight_position * mean[3],
self._std_weight_position * mean[3],
1e-2,
self._std_weight_position * mean[3],
]
std_vel = [
self._std_weight_velocity * mean[3],
self._std_weight_velocity * mean[3],
1e-5,
self._std_weight_velocity * mean[3],
]
motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
# mean = np.dot(self._motion_mat, mean)
mean = np.dot(mean, self._motion_mat.T)
covariance = (
np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T))
+ motion_cov
)
return mean, covariance
def project(self, mean, covariance):
"""Project state distribution to measurement space.
Parameters
----------
mean : ndarray
The state's mean vector (8 dimensional array).
covariance : ndarray
The state's covariance matrix (8x8 dimensional).
Returns
-------
(ndarray, ndarray)
Returns the projected mean and covariance matrix of the given state
estimate.
"""
std = [
self._std_weight_position * mean[3],
self._std_weight_position * mean[3],
1e-1,
self._std_weight_position * mean[3],
]
innovation_cov = np.diag(np.square(std))
mean = np.dot(self._update_mat, mean)
covariance = np.linalg.multi_dot(
(self._update_mat, covariance, self._update_mat.T)
)
return mean, covariance + innovation_cov
def multi_predict(self, mean, covariance):
"""Run Kalman filter prediction step (Vectorized version).
Parameters
----------
mean : ndarray
The Nx8 dimensional mean matrix of the object states at the previous
time step.
covariance : ndarray
The Nx8x8 dimensional covariance matrics of the object states at the
previous time step.
Returns
-------
(ndarray, ndarray)
Returns the mean vector and covariance matrix of the predicted
state. Unobserved velocities are initialized to 0 mean.
"""
std_pos = [
self._std_weight_position * mean[:, 3],
self._std_weight_position * mean[:, 3],
1e-2 * np.ones_like(mean[:, 3]),
self._std_weight_position * mean[:, 3],
]
std_vel = [
self._std_weight_velocity * mean[:, 3],
self._std_weight_velocity * mean[:, 3],
1e-5 * np.ones_like(mean[:, 3]),
self._std_weight_velocity * mean[:, 3],
]
sqr = np.square(np.r_[std_pos, std_vel]).T
motion_cov = []
for i in range(len(mean)):
motion_cov.append(np.diag(sqr[i]))
motion_cov = np.asarray(motion_cov)
mean = np.dot(mean, self._motion_mat.T)
left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2))
covariance = np.dot(left, self._motion_mat.T) + motion_cov
return mean, covariance
def update(self, mean, covariance, measurement):
"""Run Kalman filter correction step.
Parameters
----------
mean : ndarray
The predicted state's mean vector (8 dimensional).
covariance : ndarray
The state's covariance matrix (8x8 dimensional).
measurement : ndarray
The 4 dimensional measurement vector (x, y, a, h), where (x, y)
is the center position, a the aspect ratio, and h the height of the
bounding box.
Returns
-------
(ndarray, ndarray)
Returns the measurement-corrected state distribution.
"""
projected_mean, projected_cov = self.project(mean, covariance)
chol_factor, lower = scipy.linalg.cho_factor(
projected_cov, lower=True, check_finite=False
)
kalman_gain = scipy.linalg.cho_solve(
(chol_factor, lower),
np.dot(covariance, self._update_mat.T).T,
check_finite=False,
).T
innovation = measurement - projected_mean
new_mean = mean + np.dot(innovation, kalman_gain.T)
new_covariance = covariance - np.linalg.multi_dot(
(kalman_gain, projected_cov, kalman_gain.T)
)
return new_mean, new_covariance
def gating_distance(
self, mean, covariance, measurements, only_position=False, metric="maha"
):
"""Compute gating distance between state distribution and measurements.
A suitable distance threshold can be obtained from `chi2inv95`. If
`only_position` is False, the chi-square distribution has 4 degrees of
freedom, otherwise 2.
Parameters
----------
mean : ndarray
Mean vector over the state distribution (8 dimensional).
covariance : ndarray
Covariance of the state distribution (8x8 dimensional).
measurements : ndarray
An Nx4 dimensional matrix of N measurements, each in
format (x, y, a, h) where (x, y) is the bounding box center
position, a the aspect ratio, and h the height.
only_position : Optional[bool]
If True, distance computation is done with respect to the bounding
box center position only.
Returns
-------
ndarray
Returns an array of length N, where the i-th element contains the
squared Mahalanobis distance between (mean, covariance) and
`measurements[i]`.
"""
mean, covariance = self.project(mean, covariance)
if only_position:
mean, covariance = mean[:2], covariance[:2, :2]
measurements = measurements[:, :2]
d = measurements - mean
if metric == "gaussian":
return np.sum(d * d, axis=1)
elif metric == "maha":
cholesky_factor = np.linalg.cholesky(covariance)
z = scipy.linalg.solve_triangular(
cholesky_factor, d.T, lower=True, check_finite=False, overwrite_b=True
)
squared_maha = np.sum(z * z, axis=0)
return squared_maha
else:
raise ValueError("invalid distance metric")

237
supervision/tracker/byte_tracker/matching.py Normal file
View File

@ -0,0 +1,237 @@
import numpy as np
import scipy
from scipy.spatial.distance import cdist
from scipy.optimize import linear_sum_assignment
from supervision.tracker.byte_tracker import kalman_filter
def merge_matches(m1, m2, shape):
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(cost_matrix, indices, thresh):
matched_cost = cost_matrix[tuple(zip(*indices))]
matched_mask = matched_cost <= thresh
matches = indices[matched_mask]
unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
return matches, unmatched_a, unmatched_b
def linear_assignment(cost_matrix, thresh):
"""
Simple linear assignment
:type cost_matrix: np.ndarray
:type thresh: float
:return: matches, unmatched_a, unmatched_b
"""
if cost_matrix.size == 0:
return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
cost_matrix[cost_matrix > thresh] = thresh + 1e-4
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)
def ious(atlbrs, btlbrs):
"""
Compute cost based on IoU
:type atlbrs: list[tlbr] | np.ndarray
:type atlbrs: list[tlbr] | np.ndarray
:rtype ious np.ndarray
"""
ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
if ious.size == 0:
return ious
ious = bbox_ious(
np.ascontiguousarray(atlbrs, dtype=np.float32),
np.ascontiguousarray(btlbrs, dtype=np.float32),
)
return ious
def iou_distance(atracks, btracks):
"""
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.tlbr for track in atracks]
btlbrs = [track.tlbr for track in btracks]
_ious = ious(atlbrs, btlbrs)
cost_matrix = 1 - _ious
return cost_matrix
def v_iou_distance(atracks, btracks):
"""
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 = ious(atlbrs, btlbrs)
cost_matrix = 1 - _ious
return cost_matrix
def embedding_distance(tracks, detections, metric="cosine"):
"""
: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
)
# for i, track in enumerate(tracks):
# cost_matrix[i, :] = np.maximum(0.0, cdist(track.smooth_feat.reshape(1,-1), det_features, metric))
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, tracks, detections, only_position=False):
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, tracks, detections, only_position=False, lambda_=0.98):
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, tracks, detections):
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, detections):
if cost_matrix.size == 0:
return cost_matrix
iou_sim = 1 - cost_matrix
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 = iou_sim * det_scores
fuse_cost = 1 - fuse_sim
return fuse_cost
def bbox_ious(boxes, query_boxes):
"""
Parameters
----------
boxes: (N, 4) ndarray of float
query_boxes: (K, 4) ndarray of float
Returns
-------
overlaps: (N, K) ndarray of overlap between boxes and query_boxes
"""
N = boxes.shape[0]
K = query_boxes.shape[0]
overlaps = np.zeros((N, K), dtype=np.float32)
for k in range(K):
box_area = (query_boxes[k, 2] - query_boxes[k, 0] + 1) * (
query_boxes[k, 3] - query_boxes[k, 1] + 1
)
for n in range(N):
iw = (
min(boxes[n, 2], query_boxes[k, 2])
- max(boxes[n, 0], query_boxes[k, 0])
+ 1
)
if iw > 0:
ih = (
min(boxes[n, 3], query_boxes[k, 3])
- max(boxes[n, 1], query_boxes[k, 1])
+ 1
)
if ih > 0:
ua = float(
(boxes[n, 2] - boxes[n, 0] + 1)
* (boxes[n, 3] - boxes[n, 1] + 1)
+ box_area
- iw * ih
)
overlaps[n, k] = iw * ih / ua
return overlaps