diff --git a/README.md b/README.md index d74aada..16f0402 100755 --- a/README.md +++ b/README.md @@ -68,7 +68,13 @@ When instantiating a `DeepSort` object (as in `deepsort_tracker.py`), `polygon` - Due to special request, tensorflow embedder is available now too (very unwillingly included). - Skip nms completely in preprocessing detections if `nms_max_overlap == 1.0` (which is the default), in the original repo, nms will still be done even if threshold is set to 1.0 (probably because it was not optimised for speed). - Now able to override the `Track` class with a custom Track class (that inherits from `Track` class) for custom track logic -- Now takes in a "clock" object (see `utils/clock.py` for example), which provides date for track naming and facilities track id reset every day, preventing overflow and overly large track ids when system runs for a long time. +- Takes in today's date now, which provides date for track naming and facilities track id reset every day, preventing overflow and overly large track ids when system runs for a long time. + + ```python3 + from datetime import datetime + today = datetime.now().date() + ``` + - Now supports polygon detections. We do not track polygon points per se, but merely convert the polygon to its bounding rectangle for tracking. That said, if embedding is enabled, the embedder works on the crop around the bounding rectangle, with area not covered by the polygon masked away. [Read more here](#polygon-support). - The original `Track.to_*` methods for retrieving bounding box values returns only the Kalman predicted values. In some applications, it is better to return the bb values of the original detections the track was associated to at the current round. Added a `orig` argument which can be flagged `True` to get that. [Read more here](#getting-bounding-box-of-original-detection). - Added `get_det_supplementary` method to `Track` class, in order to pass detection related info through the track. [Read more here](#storing-supplementary-info-of-original-detection). diff --git a/deep_sort_realtime/__init__.py b/deep_sort_realtime/__init__.py index 4802e90..f901408 100755 --- a/deep_sort_realtime/__init__.py +++ b/deep_sort_realtime/__init__.py @@ -1 +1 @@ -__version__ = "1.0" +__version__ = "1.1" diff --git a/deep_sort_realtime/deep_sort/detection.py b/deep_sort_realtime/deep_sort/detection.py index d5cbc99..3a8f57c 100755 --- a/deep_sort_realtime/deep_sort/detection.py +++ b/deep_sort_realtime/deep_sort/detection.py @@ -15,7 +15,7 @@ class Detection(object): feature : array_like A feature vector that describes the object contained in this image. class_name : Optional str - Detector predicted class name. + Detector predicted class name. others : Optional any Other supplementary fields associated with detection that wants to be stored as a "memory" to be retrieve through the track downstream. @@ -31,7 +31,7 @@ class Detection(object): """ def __init__(self, ltwh, confidence, feature, class_name=None, others=None): - # def __init__(self, ltwh, feature): + # def __init__(self, ltwh, feature): self.ltwh = np.asarray(ltwh, dtype=np.float) self.confidence = float(confidence) self.feature = np.asarray(feature, dtype=np.float32) diff --git a/deep_sort_realtime/deep_sort/iou_matching.py b/deep_sort_realtime/deep_sort/iou_matching.py index 33572eb..78e3097 100755 --- a/deep_sort_realtime/deep_sort/iou_matching.py +++ b/deep_sort_realtime/deep_sort/iou_matching.py @@ -27,11 +27,15 @@ def iou(bbox, candidates): candidates_tl = candidates[:, :2] candidates_br = candidates[:, :2] + candidates[:, 2:] - tl = np.c_[np.maximum(bbox_tl[0], candidates_tl[:, 0])[:, np.newaxis], - np.maximum(bbox_tl[1], candidates_tl[:, 1])[:, np.newaxis]] - br = np.c_[np.minimum(bbox_br[0], candidates_br[:, 0])[:, np.newaxis], - np.minimum(bbox_br[1], candidates_br[:, 1])[:, np.newaxis]] - wh = np.maximum(0., br - tl) + tl = np.c_[ + np.maximum(bbox_tl[0], candidates_tl[:, 0])[:, np.newaxis], + np.maximum(bbox_tl[1], candidates_tl[:, 1])[:, np.newaxis], + ] + br = np.c_[ + np.minimum(bbox_br[0], candidates_br[:, 0])[:, np.newaxis], + np.minimum(bbox_br[1], candidates_br[:, 1])[:, np.newaxis], + ] + wh = np.maximum(0.0, br - tl) area_intersection = wh.prod(axis=1) area_bbox = bbox[2:].prod() @@ -39,8 +43,7 @@ def iou(bbox, candidates): return area_intersection / (area_bbox + area_candidates - area_intersection) -def iou_cost(tracks, detections, track_indices=None, - detection_indices=None): +def iou_cost(tracks, detections, track_indices=None, detection_indices=None): """An intersection over union distance metric. Parameters @@ -77,5 +80,5 @@ def iou_cost(tracks, detections, track_indices=None, bbox = tracks[track_idx].to_ltwh() candidates = np.asarray([detections[i].ltwh for i in detection_indices]) - cost_matrix[row, :] = 1. - iou(bbox, candidates) + cost_matrix[row, :] = 1.0 - iou(bbox, candidates) return cost_matrix diff --git a/deep_sort_realtime/deep_sort/kalman_filter.py b/deep_sort_realtime/deep_sort/kalman_filter.py index 787a76e..4e92dfd 100755 --- a/deep_sort_realtime/deep_sort/kalman_filter.py +++ b/deep_sort_realtime/deep_sort/kalman_filter.py @@ -17,7 +17,8 @@ 6: 12.592, 7: 14.067, 8: 15.507, - 9: 16.919} + 9: 16.919, +} class KalmanFilter(object): @@ -38,7 +39,7 @@ class KalmanFilter(object): """ def __init__(self): - ndim, dt = 4, 1. + ndim, dt = 4, 1.0 # Create Kalman filter model matrices. self._motion_mat = np.eye(2 * ndim, 2 * ndim) @@ -49,8 +50,8 @@ def __init__(self): # 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. / 20 - self._std_weight_velocity = 1. / 160 + self._std_weight_position = 1.0 / 20 + self._std_weight_velocity = 1.0 / 160 def initiate(self, measurement): """Create track from unassociated measurement. @@ -81,7 +82,8 @@ def initiate(self, measurement): 10 * self._std_weight_velocity * measurement[3], 10 * self._std_weight_velocity * measurement[3], 1e-5, - 10 * self._std_weight_velocity * measurement[3]] + 10 * self._std_weight_velocity * measurement[3], + ] covariance = np.diag(np.square(std)) return mean, covariance @@ -108,17 +110,21 @@ def predict(self, mean, covariance): self._std_weight_position * mean[3], self._std_weight_position * mean[3], 1e-2, - self._std_weight_position * mean[3]] + 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]] + 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) - covariance = np.linalg.multi_dot(( - self._motion_mat, covariance, self._motion_mat.T)) + motion_cov + covariance = ( + np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T)) + + motion_cov + ) return mean, covariance @@ -143,12 +149,14 @@ def project(self, mean, covariance): self._std_weight_position * mean[3], self._std_weight_position * mean[3], 1e-1, - self._std_weight_position * mean[3]] + 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)) + covariance = np.linalg.multi_dot( + (self._update_mat, covariance, self._update_mat.T) + ) return mean, covariance + innovation_cov def update(self, mean, covariance, measurement): @@ -174,19 +182,22 @@ def update(self, mean, covariance, measurement): projected_mean, projected_cov = self.project(mean, covariance) chol_factor, lower = scipy.linalg.cho_factor( - projected_cov, lower=True, check_finite=False) + 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 + (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)) + 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): + def gating_distance(self, mean, covariance, measurements, only_position=False): """Compute gating distance between state distribution and measurements. A suitable distance threshold can be obtained from `chi2inv95`. If @@ -223,7 +234,7 @@ def gating_distance(self, mean, covariance, measurements, cholesky_factor = np.linalg.cholesky(covariance) d = measurements - mean z = scipy.linalg.solve_triangular( - cholesky_factor, d.T, lower=True, check_finite=False, - overwrite_b=True) + cholesky_factor, d.T, lower=True, check_finite=False, overwrite_b=True + ) squared_maha = np.sum(z * z, axis=0) return squared_maha diff --git a/deep_sort_realtime/deep_sort/linear_assignment.py b/deep_sort_realtime/deep_sort/linear_assignment.py index fff9358..d1f6fc6 100755 --- a/deep_sort_realtime/deep_sort/linear_assignment.py +++ b/deep_sort_realtime/deep_sort/linear_assignment.py @@ -1,17 +1,23 @@ # vim: expandtab:ts=4:sw=4 from __future__ import absolute_import import numpy as np + # from sklearn.utils.linear_assignment_ import linear_assignment from scipy.optimize import linear_sum_assignment from . import kalman_filter -INFTY_COST = 1e+5 +INFTY_COST = 1e5 def min_cost_matching( - distance_metric, max_distance, tracks, detections, track_indices=None, - detection_indices=None): + distance_metric, + max_distance, + tracks, + detections, + track_indices=None, + detection_indices=None, +): """Solve linear assignment problem. Parameters @@ -53,8 +59,7 @@ def min_cost_matching( if len(detection_indices) == 0 or len(track_indices) == 0: return [], track_indices, detection_indices # Nothing to match. - cost_matrix = distance_metric( - tracks, detections, track_indices, detection_indices) + cost_matrix = distance_metric(tracks, detections, track_indices, detection_indices) cost_matrix[cost_matrix > max_distance] = max_distance + 1e-5 # indices = linear_assignment(cost_matrix) indices = np.vstack(linear_sum_assignment(cost_matrix)).T @@ -78,8 +83,14 @@ def min_cost_matching( def matching_cascade( - distance_metric, max_distance, cascade_depth, tracks, detections, - track_indices=None, detection_indices=None): + distance_metric, + max_distance, + cascade_depth, + tracks, + detections, + track_indices=None, + detection_indices=None, +): """Run matching cascade. Parameters @@ -128,24 +139,34 @@ def matching_cascade( break track_indices_l = [ - k for k in track_indices - if tracks[k].time_since_update == 1 + level + k for k in track_indices if tracks[k].time_since_update == 1 + level ] if len(track_indices_l) == 0: # Nothing to match at this level continue - matches_l, _, unmatched_detections = \ - min_cost_matching( - distance_metric, max_distance, tracks, detections, - track_indices_l, unmatched_detections) + matches_l, _, unmatched_detections = min_cost_matching( + distance_metric, + max_distance, + tracks, + detections, + track_indices_l, + unmatched_detections, + ) matches += matches_l unmatched_tracks = list(set(track_indices) - set(k for k, _ in matches)) return matches, unmatched_tracks, unmatched_detections def gate_cost_matrix( - kf, cost_matrix, tracks, detections, track_indices, detection_indices, - gated_cost=INFTY_COST, only_position=False): + kf, + cost_matrix, + tracks, + detections, + track_indices, + detection_indices, + gated_cost=INFTY_COST, + only_position=False, +): """Invalidate infeasible entries in cost matrix based on the state distributions obtained by Kalman filtering. @@ -182,11 +203,11 @@ def gate_cost_matrix( """ gating_dim = 2 if only_position else 4 gating_threshold = kalman_filter.chi2inv95[gating_dim] - measurements = np.asarray( - [detections[i].to_xyah() for i in detection_indices]) + measurements = np.asarray([detections[i].to_xyah() for i in detection_indices]) for row, track_idx in enumerate(track_indices): track = tracks[track_idx] gating_distance = kf.gating_distance( - track.mean, track.covariance, measurements, only_position) + track.mean, track.covariance, measurements, only_position + ) cost_matrix[row, gating_distance > gating_threshold] = gated_cost return cost_matrix diff --git a/deep_sort_realtime/deep_sort/nn_matching.py b/deep_sort_realtime/deep_sort/nn_matching.py index 2e7bfea..df6445e 100755 --- a/deep_sort_realtime/deep_sort/nn_matching.py +++ b/deep_sort_realtime/deep_sort/nn_matching.py @@ -23,8 +23,8 @@ def _pdist(a, b): if len(a) == 0 or len(b) == 0: return np.zeros((len(a), len(b))) a2, b2 = np.square(a).sum(axis=1), np.square(b).sum(axis=1) - r2 = -2. * np.dot(a, b.T) + a2[:, None] + b2[None, :] - r2 = np.clip(r2, 0., float(np.inf)) + r2 = -2.0 * np.dot(a, b.T) + a2[:, None] + b2[None, :] + r2 = np.clip(r2, 0.0, float(np.inf)) return r2 @@ -51,11 +51,11 @@ def _cosine_distance(a, b, data_is_normalized=False): if not data_is_normalized: a = np.asarray(a) / np.linalg.norm(a, axis=1, keepdims=True) b = np.asarray(b) / np.linalg.norm(b, axis=1, keepdims=True) - return 1. - np.dot(a, b.T) + return 1.0 - np.dot(a, b.T) def _nn_euclidean_distance(x, y): - """ Helper function for nearest neighbor distance metric (Euclidean). + """Helper function for nearest neighbor distance metric (Euclidean). Parameters ---------- @@ -76,7 +76,7 @@ def _nn_euclidean_distance(x, y): def _nn_cosine_distance(x, y): - """ Helper function for nearest neighbor distance metric (cosine). + """Helper function for nearest neighbor distance metric (cosine). Parameters ---------- @@ -122,14 +122,12 @@ class NearestNeighborDistanceMetric(object): def __init__(self, metric, matching_threshold, budget=None): - if metric == "euclidean": self._metric = _nn_euclidean_distance elif metric == "cosine": self._metric = _nn_cosine_distance else: - raise ValueError( - "Invalid metric; must be either 'euclidean' or 'cosine'") + raise ValueError("Invalid metric; must be either 'euclidean' or 'cosine'") self.matching_threshold = matching_threshold self.budget = budget self.samples = {} @@ -150,7 +148,7 @@ def partial_fit(self, features, targets, active_targets): for feature, target in zip(features, targets): self.samples.setdefault(target, []).append(feature) if self.budget is not None: - self.samples[target] = self.samples[target][-self.budget:] + self.samples[target] = self.samples[target][-self.budget :] self.samples = {k: self.samples[k] for k in active_targets} def distance(self, features, targets): diff --git a/deep_sort_realtime/deep_sort/track.py b/deep_sort_realtime/deep_sort/track.py index 7fd9dc2..f556b43 100755 --- a/deep_sort_realtime/deep_sort/track.py +++ b/deep_sort_realtime/deep_sort/track.py @@ -69,8 +69,19 @@ class Track: """ - def __init__(self, mean, covariance, track_id, n_init, max_age, - feature=None, original_ltwh=None, det_class=None, det_conf=None, others=None): + def __init__( + self, + mean, + covariance, + track_id, + n_init, + max_age, + feature=None, + original_ltwh=None, + det_class=None, + det_conf=None, + others=None, + ): self.mean = mean self.covariance = covariance self.track_id = track_id @@ -159,21 +170,21 @@ def to_ltrb(self, orig=False): return ret def get_det_conf(self): - ''' + """ `det_conf` will be None is there are no associated detection this round - ''' + """ return self.det_conf - + def get_det_class(self): - ''' + """ Only `det_class` will be persisted in the track even if there are no associated detection this round. - ''' + """ return self.det_class def get_det_supplementary(self): - ''' + """ Get supplementary info associated with the detection. Will be None is there are no associated detection this round. - ''' + """ return self.others def predict(self, kf): @@ -207,7 +218,8 @@ def update(self, kf, detection): """ self.original_ltwh = detection.get_ltwh() self.mean, self.covariance = kf.update( - self.mean, self.covariance, detection.to_xyah()) + self.mean, self.covariance, detection.to_xyah() + ) self.features.append(detection.feature) self.det_conf = detection.confidence self.det_class = detection.class_name @@ -218,19 +230,16 @@ def update(self, kf, detection): self.time_since_update = 0 if self.state == TrackState.Tentative and self.hits >= self._n_init: self.state = TrackState.Confirmed - def mark_missed(self): - """Mark this track as missed (no association at the current time step). - """ + """Mark this track as missed (no association at the current time step).""" if self.state == TrackState.Tentative: self.state = TrackState.Deleted elif self.time_since_update > self._max_age: self.state = TrackState.Deleted def is_tentative(self): - """Returns True if this track is tentative (unconfirmed). - """ + """Returns True if this track is tentative (unconfirmed).""" return self.state == TrackState.Tentative def is_confirmed(self): diff --git a/deep_sort_realtime/deep_sort/tracker.py b/deep_sort_realtime/deep_sort/tracker.py index 712f19b..1e8aa24 100755 --- a/deep_sort_realtime/deep_sort/tracker.py +++ b/deep_sort_realtime/deep_sort/tracker.py @@ -1,5 +1,6 @@ # vim: expandtab:ts=4:sw=4 from __future__ import absolute_import +from datetime import datetime import numpy as np from . import kalman_filter from . import linear_assignment @@ -21,6 +22,8 @@ class Tracker: Number of consecutive detections before the track is confirmed. The track state is set to `Deleted` if a miss occurs within the first `n_init` frames. + today: Optional[datetime.date] + Provide today's date, for naming of tracks Attributes ---------- @@ -37,11 +40,16 @@ class Tracker: """ - def __init__(self, metric, max_iou_distance=0.7, max_age=30, n_init=3, override_track_class=None, clock=None): - # assert clock is not None - self.clock = clock - if self.clock: - self.today = self.clock.get_now_SGT().date() + def __init__( + self, + metric, + max_iou_distance=0.7, + max_age=30, + n_init=3, + override_track_class=None, + today=None, + ): + self.today = today self.metric = metric self.max_iou_distance = max_iou_distance self.max_age = max_age @@ -55,7 +63,6 @@ def __init__(self, metric, max_iou_distance=0.7, max_age=30, n_init=3, override_ self.track_class = override_track_class else: self.track_class = Track - def predict(self): """Propagate track state distributions one time step forward. @@ -65,29 +72,30 @@ def predict(self): for track in self.tracks: track.predict(self.kf) - def update(self, detections): + def update(self, detections, today=None): """Perform measurement update and track management. Parameters ---------- detections : List[deep_sort.detection.Detection] A list of detections at the current time step. - + today: Optional[datetime.date] + Provide today's date, for naming of tracks """ - if self.clock: + if self.today: + if today is None: + today = datetime.now().date() # Check if its a new day, then refresh idx - if self.clock.get_now_SGT().date() != self.today: - self.today = self.clock.get_now_SGT().date() + if today != self.today: + self.today = today self._next_id = 1 # Run matching cascade. - matches, unmatched_tracks, unmatched_detections = \ - self._match(detections) + matches, unmatched_tracks, unmatched_detections = self._match(detections) # Update track set. for track_idx, detection_idx in matches: - self.tracks[track_idx].update( - self.kf, detections[detection_idx]) + self.tracks[track_idx].update(self.kf, detections[detection_idx]) for track_idx in unmatched_tracks: self.tracks[track_idx].mark_missed() for detection_idx in unmatched_detections: @@ -112,43 +120,59 @@ def update(self, detections): targets += [track.track_id for _ in track.features] track.features = [track.features[-1]] self.metric.partial_fit( - np.asarray(features), np.asarray(targets), active_targets) + np.asarray(features), np.asarray(targets), active_targets + ) def _match(self, detections): - def gated_metric(tracks, dets, track_indices, detection_indices): features = np.array([dets[i].feature for i in detection_indices]) targets = np.array([tracks[i].track_id for i in track_indices]) cost_matrix = self.metric.distance(features, targets) cost_matrix = linear_assignment.gate_cost_matrix( - self.kf, cost_matrix, tracks, dets, track_indices, - detection_indices) + self.kf, cost_matrix, tracks, dets, track_indices, detection_indices + ) return cost_matrix # Split track set into confirmed and unconfirmed tracks. - confirmed_tracks = [ - i for i, t in enumerate(self.tracks) if t.is_confirmed()] + confirmed_tracks = [i for i, t in enumerate(self.tracks) if t.is_confirmed()] unconfirmed_tracks = [ - i for i, t in enumerate(self.tracks) if not t.is_confirmed()] + i for i, t in enumerate(self.tracks) if not t.is_confirmed() + ] # Associate confirmed tracks using appearance features. - matches_a, unmatched_tracks_a, unmatched_detections = \ - linear_assignment.matching_cascade( - gated_metric, self.metric.matching_threshold, self.max_age, - self.tracks, detections, confirmed_tracks) + ( + matches_a, + unmatched_tracks_a, + unmatched_detections, + ) = linear_assignment.matching_cascade( + gated_metric, + self.metric.matching_threshold, + self.max_age, + self.tracks, + detections, + confirmed_tracks, + ) # Associate remaining tracks together with unconfirmed tracks using IOU. iou_track_candidates = unconfirmed_tracks + [ - k for k in unmatched_tracks_a if - self.tracks[k].time_since_update == 1] + k for k in unmatched_tracks_a if self.tracks[k].time_since_update == 1 + ] unmatched_tracks_a = [ - k for k in unmatched_tracks_a if - self.tracks[k].time_since_update != 1] - matches_b, unmatched_tracks_b, unmatched_detections = \ - linear_assignment.min_cost_matching( - iou_matching.iou_cost, self.max_iou_distance, self.tracks, - detections, iou_track_candidates, unmatched_detections) + k for k in unmatched_tracks_a if self.tracks[k].time_since_update != 1 + ] + ( + matches_b, + unmatched_tracks_b, + unmatched_detections, + ) = linear_assignment.min_cost_matching( + iou_matching.iou_cost, + self.max_iou_distance, + self.tracks, + detections, + iou_track_candidates, + unmatched_detections, + ) matches = matches_a + matches_b unmatched_tracks = list(set(unmatched_tracks_a + unmatched_tracks_b)) @@ -156,13 +180,22 @@ def gated_metric(tracks, dets, track_indices, detection_indices): def _initiate_track(self, detection): mean, covariance = self.kf.initiate(detection.to_xyah()) - - if self.clock: - track_id = '{}_{}'.format(self.clock.get_now_SGT_date_str(), self._next_id) + + if self.today: + track_id = "{}_{}".format(self.today, self._next_id) else: - track_id = '{}'.format(self._next_id) - self.tracks.append(self.track_class( - mean, covariance, track_id, self.n_init, self.max_age, - # mean, covariance, self._next_id, self.n_init, self.max_age, - feature=detection.feature, det_class=detection.class_name, det_conf=detection.confidence)) + track_id = "{}".format(self._next_id) + self.tracks.append( + self.track_class( + mean, + covariance, + track_id, + self.n_init, + self.max_age, + # mean, covariance, self._next_id, self.n_init, self.max_age, + feature=detection.feature, + det_class=detection.class_name, + det_conf=detection.confidence, + ) + ) self._next_id += 1 diff --git a/deep_sort_realtime/deepsort_tracker.py b/deep_sort_realtime/deepsort_tracker.py index b8ea1bc..2378ce3 100644 --- a/deep_sort_realtime/deepsort_tracker.py +++ b/deep_sort_realtime/deepsort_tracker.py @@ -10,11 +10,24 @@ from deep_sort_realtime.utils.nms import non_max_suppression logger = logging.getLogger(__name__) + + class DeepSort(object): + def __init__( + self, + max_age=30, + nms_max_overlap=1.0, + max_cosine_distance=0.2, + nn_budget=None, + override_track_class=None, + embedder=True, + half=True, + bgr=True, + polygon=False, + today=None, + ): + """ - def __init__(self, max_age = 30, nms_max_overlap=1.0, max_cosine_distance=0.2, nn_budget=None, override_track_class=None, clock=None, embedder=True, half=True, bgr=True, polygon=False): - ''' - Parameters ---------- max_age : Optional[int] = 30 @@ -27,8 +40,6 @@ def __init__(self, max_age = 30, nms_max_overlap=1.0, max_cosine_distance=0.2, n Maximum size of the appearance descriptors, if None, no budget is enforced override_track_class : Optional[object] = None Giving this will override default Track class, this must inherit Track - clock : Optional[object] = None - Clock custom object provides date for track naming and facilitates track id reset every day, preventing overflow and overly large track ids. For example clock class, please see `utils/clock.py` embedder : Optional[bool] = True Whether to use in-built embedder or not. If False, then embeddings must be given during update half : Optional[bool] = True @@ -37,30 +48,43 @@ def __init__(self, max_age = 30, nms_max_overlap=1.0, max_cosine_distance=0.2, n Whether frame given to embedder is expected to be BGR or not (RGB) polygon: Optional[bool] = False Whether detections are polygons (e.g. oriented bounding boxes) - ''' - # self.video_info = video_info - # assert clock is not None + today: Optional[datetime.date] + Provide today's date, for naming of tracks + """ self.nms_max_overlap = nms_max_overlap metric = nn_matching.NearestNeighborDistanceMetric( - "cosine", max_cosine_distance, nn_budget) - self.tracker = Tracker(metric, max_age = max_age, override_track_class=override_track_class, clock=clock) + "cosine", max_cosine_distance, nn_budget + ) + self.tracker = Tracker( + metric, + max_age=max_age, + override_track_class=override_track_class, + today=today, + ) if embedder: - from deep_sort_realtime.embedder.embedder_pytorch import MobileNetv2_Embedder as Embedder + from deep_sort_realtime.embedder.embedder_pytorch import ( + MobileNetv2_Embedder as Embedder, + ) + self.embedder = Embedder(half=half, max_batch_size=16, bgr=bgr) else: self.embedder = None self.polygon = polygon - logger.info('DeepSort Tracker initialised') - logger.info(f'- max age: {max_age}') - logger.info(f'- appearance threshold: {max_cosine_distance}') - logger.info(f'- nms threshold: {"OFF" if self.nms_max_overlap==1.0 else self.nms_max_overlap }') - logger.info(f'- max num of appearance features: {nn_budget}') - logger.info(f'- overriding track class : {"No" if override_track_class is None else "Yes"}' ) - logger.info(f'- clock : {"No" if clock is None else "Yes"}' ) - logger.info(f'- in-build embedder : {"No" if self.embedder is None else "Yes"}' ) - logger.info(f'- polygon detections : {"No" if polygon is False else "Yes"}' ) - - def update_tracks(self, raw_detections, embeds=None, frame=None): + logger.info("DeepSort Tracker initialised") + logger.info(f"- max age: {max_age}") + logger.info(f"- appearance threshold: {max_cosine_distance}") + logger.info( + f'- nms threshold: {"OFF" if self.nms_max_overlap==1.0 else self.nms_max_overlap }' + ) + logger.info(f"- max num of appearance features: {nn_budget}") + logger.info( + f'- overriding track class : {"No" if override_track_class is None else "Yes"}' + ) + logger.info(f'- today given : {"No" if today is None else "Yes"}') + logger.info(f'- in-build embedder : {"No" if self.embedder is None else "Yes"}') + logger.info(f'- polygon detections : {"No" if polygon is False else "Yes"}') + + def update_tracks(self, raw_detections, embeds=None, frame=None, today=None): """Run multi-target tracker on a particular sequence. @@ -69,11 +93,13 @@ def update_tracks(self, raw_detections, embeds=None, frame=None): raw_detections (horizontal bb) : List[ Tuple[ List[float or int], float, str ] ] List of detections, each in tuples of ( [left,top,w,h] , confidence, detection_class) raw_detections (polygon) : List[ List[float], List[int or str], List[float] ] - List of Polygons, Classes, Confidences. All 3 sublists of the same length. A polygon defined as a ndarray-like [x1,y1,x2,y2,...]. + List of Polygons, Classes, Confidences. All 3 sublists of the same length. A polygon defined as a ndarray-like [x1,y1,x2,y2,...]. embeds : Optional[ List[] ] = None List of appearance features corresponding to detections frame : Optional [ np.ndarray ] = None if embeds not given, Image frame must be given here, in [H,W,C]. + today: Optional[datetime.date] + Provide today's date, for naming of tracks Returns ------- @@ -83,44 +109,47 @@ def update_tracks(self, raw_detections, embeds=None, frame=None): if embeds is None: if self.embedder is None: - raise Exception('Embedder not created during init so embeddings must be given now!') + raise Exception( + "Embedder not created during init so embeddings must be given now!" + ) if frame is None: - raise Exception('either embeddings or frame must be given!') + raise Exception("either embeddings or frame must be given!") if not self.polygon: - raw_detections = [ d for d in raw_detections if d[0][2] > 0 and d[0][3] > 0] + raw_detections = [d for d in raw_detections if d[0][2] > 0 and d[0][3] > 0] if embeds is None: embeds = self.generate_embeds(frame, raw_detections) - + # Proper deep sort detection objects that consist of bbox, confidence and embedding. detections = self.create_detections(raw_detections, embeds) else: polygons, bounding_rects = self.process_polygons(raw_detections[0]) - + if embeds is None: embeds = self.generate_embeds_poly(frame, polygons, bounding_rects) - + # Proper deep sort detection objects that consist of bbox, confidence and embedding. - detections = self.create_detections_poly(raw_detections, embeds, bounding_rects) + detections = self.create_detections_poly( + raw_detections, embeds, bounding_rects + ) # Run non-maxima suppression. boxes = np.array([d.ltwh for d in detections]) scores = np.array([d.confidence for d in detections]) if self.nms_max_overlap < 1.0: # nms_tic = time.perf_counter() - indices = non_max_suppression( - boxes, self.nms_max_overlap, scores) + indices = non_max_suppression(boxes, self.nms_max_overlap, scores) # nms_toc = time.perf_counter() # logger.debug(f'nms time: {nms_toc-nms_tic}s') detections = [detections[i] for i in indices] # Update tracker. self.tracker.predict() - self.tracker.update(detections) + self.tracker.update(detections, today=today) return self.tracker.tracks - + def refresh_track_ids(self): self.tracker._next_id @@ -134,25 +163,34 @@ def generate_embeds_poly(self, frame, polygons, bounding_rects): def create_detections(self, raw_dets, embeds): detection_list = [] - for raw_det, embed in zip(raw_dets,embeds): - detection_list.append(Detection(raw_det[0], raw_det[1], embed, class_name=raw_det[2])) #raw_det = [bbox, conf_score, class] + for raw_det, embed in zip(raw_dets, embeds): + detection_list.append( + Detection(raw_det[0], raw_det[1], embed, class_name=raw_det[2]) + ) # raw_det = [bbox, conf_score, class] return detection_list def create_detections_poly(self, dets, embeds, bounding_rects): detection_list = [] dets.extend([embeds, bounding_rects]) for raw_polygon, cl, score, embed, bounding_rect in zip(*dets): - x,y,w,h = bounding_rect + x, y, w, h = bounding_rect x = max(0, x) y = max(0, y) - bbox = [x,y,w,h] - detection_list.append(Detection(bbox, score, embed, class_name=cl, others=raw_polygon)) + bbox = [x, y, w, h] + detection_list.append( + Detection(bbox, score, embed, class_name=cl, others=raw_polygon) + ) return detection_list @staticmethod def process_polygons(raw_polygons): - polygons = [ [ polygon[x:x+2] for x in range(0, len(polygon), 2) ]for polygon in raw_polygons ] - bounding_rects = [ cv2.boundingRect(np.array([polygon]).astype(int)) for polygon in polygons ] + polygons = [ + [polygon[x : x + 2] for x in range(0, len(polygon), 2)] + for polygon in raw_polygons + ] + bounding_rects = [ + cv2.boundingRect(np.array([polygon]).astype(int)) for polygon in polygons + ] return polygons, bounding_rects @staticmethod @@ -160,7 +198,7 @@ def crop_bb(frame, raw_dets): crops = [] im_height, im_width = frame.shape[:2] for detection in raw_dets: - l,t,w,h = [int(x) for x in detection[0]] + l, t, w, h = [int(x) for x in detection[0]] r = l + w b = t + h crop_l = max(0, l) @@ -169,7 +207,7 @@ def crop_bb(frame, raw_dets): crop_b = min(im_height, b) crops.append(frame[crop_t:crop_b, crop_l:crop_r]) return crops - + @staticmethod def crop_poly_pad_black(frame, polygons, bounding_rects): masked_polys = [] @@ -177,17 +215,17 @@ def crop_poly_pad_black(frame, polygons, bounding_rects): for polygon, bounding_rect in zip(polygons, bounding_rects): mask = np.zeros(frame.shape, dtype=np.uint8) polygon_mask = np.array([polygon]).astype(int) - cv2.fillPoly(mask, polygon_mask, color=(255,255,255)) + cv2.fillPoly(mask, polygon_mask, color=(255, 255, 255)) # apply the mask masked_image = cv2.bitwise_and(frame, mask) # crop masked image - x,y,w,h = bounding_rect + x, y, w, h = bounding_rect crop_l = max(0, x) - crop_r = min(im_width, x+w) + crop_r = min(im_width, x + w) crop_t = max(0, y) - crop_b = min(im_height, y+h) + crop_b = min(im_height, y + h) cropped = masked_image[crop_t:crop_b, crop_l:crop_r].copy() masked_polys.append(np.array(cropped)) return masked_polys diff --git a/deep_sort_realtime/embedder/embedder_pytorch.py b/deep_sort_realtime/embedder/embedder_pytorch.py index 3da4b5b..590ff0c 100644 --- a/deep_sort_realtime/embedder/embedder_pytorch.py +++ b/deep_sort_realtime/embedder/embedder_pytorch.py @@ -1,4 +1,4 @@ -import os +import os import logging import cv2 @@ -11,17 +11,21 @@ logger = logging.getLogger(__name__) -MOBILENETV2_BOTTLENECK_WTS = pkg_resources.resource_filename('deep_sort_realtime', 'embedder/weights/mobilenetv2_bottleneck_wts.pt') +MOBILENETV2_BOTTLENECK_WTS = pkg_resources.resource_filename( + "deep_sort_realtime", "embedder/weights/mobilenetv2_bottleneck_wts.pt" +) INPUT_WIDTH = 224 + def batch(iterable, bs=1): l = len(iterable) for ndx in range(0, l, bs): - yield iterable[ndx:min(ndx + bs, l)] + yield iterable[ndx : min(ndx + bs, l)] + class MobileNetv2_Embedder(object): - ''' - MobileNetv2_Embedder loads a Mobilenetv2 pretrained on Imagenet1000, with classification layer removed, exposing the bottleneck layer, outputing a feature of size 1280. + """ + MobileNetv2_Embedder loads a Mobilenetv2 pretrained on Imagenet1000, with classification layer removed, exposing the bottleneck layer, outputing a feature of size 1280. Params ------ @@ -30,15 +34,18 @@ class MobileNetv2_Embedder(object): - max_batch_size (optional, int) : max batch size for embedder, defaults to 16 - bgr (optional, Bool) : boolean flag indicating if input frames are bgr or not, defaults to True - ''' - def __init__(self, model_wts_path = None, half=True, max_batch_size = 16, bgr=True): + """ + + def __init__(self, model_wts_path=None, half=True, max_batch_size=16, bgr=True): if model_wts_path is None: model_wts_path = MOBILENETV2_BOTTLENECK_WTS - assert os.path.exists(model_wts_path),f'Mobilenetv2 model path {model_wts_path} does not exists!' - self.model = MobileNetV2_bottle(input_size=INPUT_WIDTH, width_mult=1.) + assert os.path.exists( + model_wts_path + ), f"Mobilenetv2 model path {model_wts_path} does not exists!" + self.model = MobileNetV2_bottle(input_size=INPUT_WIDTH, width_mult=1.0) self.model.load_state_dict(torch.load(model_wts_path)) - self.model.cuda() #loads model to gpu - self.model.eval() #inference mode, deactivates dropout layers + self.model.cuda() # loads model to gpu + self.model.eval() # inference mode, deactivates dropout layers self.max_batch_size = max_batch_size self.bgr = bgr @@ -46,16 +53,16 @@ def __init__(self, model_wts_path = None, half=True, max_batch_size = 16, bgr=Tr self.half = half if self.half: self.model.half() - logger.info('MobileNetV2 Embedder for Deep Sort initialised') - logger.info(f'- half precision: {self.half}') - logger.info(f'- max batch size: {self.max_batch_size}') - logger.info(f'- expects BGR: {self.bgr}') + logger.info("MobileNetV2 Embedder for Deep Sort initialised") + logger.info(f"- half precision: {self.half}") + logger.info(f"- max batch size: {self.max_batch_size}") + logger.info(f"- expects BGR: {self.bgr}") zeros = np.zeros((100, 100, 3), dtype=np.uint8) - self.predict([zeros]) #warmup + self.predict([zeros]) # warmup def preprocess(self, np_image): - ''' + """ Preprocessing for embedder network: Flips BGR to RGB, resize, convert to torch tensor, normalise with imagenet mean and variance, reshape. Note: input image yet to be loaded to GPU through tensor.cuda() Parameters @@ -67,39 +74,43 @@ def preprocess(self, np_image): ------- Torch Tensor - ''' + """ if self.bgr: - np_image_rgb = np_image[...,::-1] + np_image_rgb = np_image[..., ::-1] else: np_image_rgb = np_image input_image = cv2.resize(np_image_rgb, (INPUT_WIDTH, INPUT_WIDTH)) - trans = transforms.Compose([ - transforms.ToTensor(), - transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - ]) + trans = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), + ] + ) input_image = trans(input_image) - input_image = input_image.view(1,3,INPUT_WIDTH,INPUT_WIDTH) - + input_image = input_image.view(1, 3, INPUT_WIDTH, INPUT_WIDTH) + return input_image def predict(self, np_images): - ''' + """ batch inference Params ------ np_images : list of ndarray list of (H x W x C), bgr or rgb according to self.bgr - + Returns ------ list of features (np.array with dim = 1280) - ''' + """ all_feats = [] - preproc_imgs = [ self.preprocess(img) for img in np_images ] + preproc_imgs = [self.preprocess(img) for img in np_images] for this_batch in batch(preproc_imgs, bs=self.max_batch_size): this_batch = torch.cat(this_batch, dim=0) @@ -107,7 +118,7 @@ def predict(self, np_images): if self.half: this_batch = this_batch.half() output = self.model.forward(this_batch) - + all_feats.extend(output.cpu().data.numpy()) return all_feats diff --git a/deep_sort_realtime/embedder/embedder_tf.py b/deep_sort_realtime/embedder/embedder_tf.py index efd7b4b..66684a9 100644 --- a/deep_sort_realtime/embedder/embedder_tf.py +++ b/deep_sort_realtime/embedder/embedder_tf.py @@ -1,4 +1,4 @@ -import os +import os import logging from pathlib import Path @@ -7,11 +7,14 @@ import pkg_resources import tensorflow as tf -MOBILENETV2_BOTTLENECK_WTS = pkg_resources.resource_filename('deep_sort_realtime', 'embedder/weights/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224.h5') +MOBILENETV2_BOTTLENECK_WTS = pkg_resources.resource_filename( + "deep_sort_realtime", + "embedder/weights/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224.h5", +) logger = logging.getLogger(__name__) -gpus = tf.config.experimental.list_physical_devices('GPU') +gpus = tf.config.experimental.list_physical_devices("GPU") if gpus: # Currently, memory growth needs to be the same across GPUs for gpu in gpus: @@ -19,20 +22,22 @@ INPUT_WIDTH = 224 + def batch(iterable, bs=1): l = len(iterable) for ndx in range(0, l, bs): - yield iterable[ndx:min(ndx + bs, l)] + yield iterable[ndx : min(ndx + bs, l)] + -def get_mobilenetv2_with_preproc(wts='imagenet'): - i = tf.keras.layers.Input([None, None, 3], dtype = tf.uint8) +def get_mobilenetv2_with_preproc(wts="imagenet"): + i = tf.keras.layers.Input([None, None, 3], dtype=tf.uint8) x = tf.cast(i, tf.float32) x = tf.keras.applications.mobilenet_v2.preprocess_input(x) full_model = tf.keras.applications.mobilenet_v2.MobileNetV2( - input_shape=None, + input_shape=None, weights=str(wts), - classifier_activation=None, + classifier_activation=None, ) core_model = tf.keras.Model(full_model.input, full_model.layers[-2].output) @@ -42,9 +47,10 @@ def get_mobilenetv2_with_preproc(wts='imagenet'): model.summary() return model + class MobileNetv2_Embedder(object): - ''' - MobileNetv2_Embedder loads a Mobilenetv2 pretrained on Imagenet1000, with classification layer removed, exposing the bottleneck layer, outputing a feature of size 1280. + """ + MobileNetv2_Embedder loads a Mobilenetv2 pretrained on Imagenet1000, with classification layer removed, exposing the bottleneck layer, outputing a feature of size 1280. Params ------ @@ -52,27 +58,30 @@ class MobileNetv2_Embedder(object): - max_batch_size (optional, int) : max batch size for embedder, defaults to 16 - bgr (optional, Bool) : boolean flag indicating if input frames are bgr or not, defaults to True - ''' - def __init__(self, model_wts_path = None, max_batch_size = 16, bgr=True): + """ + + def __init__(self, model_wts_path=None, max_batch_size=16, bgr=True): if model_wts_path is None: model_wts_path = MOBILENETV2_BOTTLENECK_WTS model_wts_path = Path(model_wts_path) - assert model_wts_path.is_file(),f'Mobilenetv2 model path {model_wts_path} does not exists!' + assert ( + model_wts_path.is_file() + ), f"Mobilenetv2 model path {model_wts_path} does not exists!" self.model = get_mobilenetv2_with_preproc(wts=model_wts_path) self.max_batch_size = max_batch_size self.bgr = bgr - logger.info('MobileNetV2 Embedder (tf) for Deep Sort initialised') - logger.info(f'- max batch size: {self.max_batch_size}') - logger.info(f'- expects BGR: {self.bgr}') + logger.info("MobileNetV2 Embedder (tf) for Deep Sort initialised") + logger.info(f"- max batch size: {self.max_batch_size}") + logger.info(f"- expects BGR: {self.bgr}") zeros = np.zeros((100, 100, 3), dtype=np.uint8) - self.predict([zeros, zeros]) #warmup + self.predict([zeros, zeros]) # warmup def preprocess(self, np_image): - ''' + """ Parameters ---------- np_image : ndarray @@ -82,31 +91,31 @@ def preprocess(self, np_image): ------- TF Tensor - ''' + """ if self.bgr: - np_image_rgb = np_image[...,::-1] + np_image_rgb = np_image[..., ::-1] else: np_image_rgb = np_image np_image_rgb = cv2.resize(np_image_rgb, (INPUT_WIDTH, INPUT_WIDTH)) return tf.convert_to_tensor(np_image_rgb) def predict(self, np_images): - ''' + """ batch inference Params ------ np_images : list of ndarray list of (H x W x C), bgr or rgb according to self.bgr - + Returns ------ list of features (np.array with dim = 1280) - ''' + """ all_feats = [] - preproc_imgs = [ self.preprocess(img) for img in np_images ] + preproc_imgs = [self.preprocess(img) for img in np_images] for this_batch in batch(preproc_imgs, bs=self.max_batch_size): this_batch = tf.stack(this_batch, axis=0) diff --git a/deep_sort_realtime/embedder/mobilenetv2_bottle.py b/deep_sort_realtime/embedder/mobilenetv2_bottle.py index 9a9ee72..a825117 100644 --- a/deep_sort_realtime/embedder/mobilenetv2_bottle.py +++ b/deep_sort_realtime/embedder/mobilenetv2_bottle.py @@ -1,11 +1,12 @@ import torch.nn as nn import math + def conv_bn(inp, oup, stride): return nn.Sequential( nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup), - nn.ReLU6(inplace=True) + nn.ReLU6(inplace=True), ) @@ -13,9 +14,10 @@ def conv_1x1_bn(inp, oup): return nn.Sequential( nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup), - nn.ReLU6(inplace=True) + nn.ReLU6(inplace=True), ) + class InvertedResidual(nn.Module): def __init__(self, inp, oup, stride, expand_ratio): super(InvertedResidual, self).__init__() @@ -28,7 +30,9 @@ def __init__(self, inp, oup, stride, expand_ratio): if expand_ratio == 1: self.conv = nn.Sequential( # dw - nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False), + nn.Conv2d( + hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False + ), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True), # pw-linear @@ -42,7 +46,9 @@ def __init__(self, inp, oup, stride, expand_ratio): nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True), # dw - nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False), + nn.Conv2d( + hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False + ), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True), # pw-linear @@ -58,7 +64,7 @@ def forward(self, x): class MobileNetV2_bottle(nn.Module): - def __init__(self, input_size=224, width_mult=1.): + def __init__(self, input_size=224, width_mult=1.0): super(MobileNetV2_bottle, self).__init__() block = InvertedResidual input_channel = 32 @@ -77,16 +83,22 @@ def __init__(self, input_size=224, width_mult=1.): # building first layer assert input_size % 32 == 0 input_channel = int(input_channel * width_mult) - self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel + self.last_channel = ( + int(last_channel * width_mult) if width_mult > 1.0 else last_channel + ) self.features = [conv_bn(3, input_channel, 2)] # building inverted residual blocks for t, c, n, s in interverted_residual_setting: output_channel = int(c * width_mult) for i in range(n): if i == 0: - self.features.append(block(input_channel, output_channel, s, expand_ratio=t)) + self.features.append( + block(input_channel, output_channel, s, expand_ratio=t) + ) else: - self.features.append(block(input_channel, output_channel, 1, expand_ratio=t)) + self.features.append( + block(input_channel, output_channel, 1, expand_ratio=t) + ) input_channel = output_channel # building last several layers self.features.append(conv_1x1_bn(input_channel, self.last_channel)) @@ -111,7 +123,7 @@ def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels - m.weight.data.normal_(0, math.sqrt(2. / n)) + m.weight.data.normal_(0, math.sqrt(2.0 / n)) if m.bias is not None: m.bias.data.zero_() elif isinstance(m, nn.BatchNorm2d): diff --git a/deep_sort_realtime/utils/clock.py b/deep_sort_realtime/utils/clock.py deleted file mode 100644 index 09d3be0..0000000 --- a/deep_sort_realtime/utils/clock.py +++ /dev/null @@ -1,42 +0,0 @@ -import os -import logging - -from datetime import datetime, timedelta - -logger = logging.getLogger(__name__) -class Clock(object): - def __init__(self): - tz_offset_hr = float(os.environ.get('TIMEZONE_OFFSET','0.0')) - self.tz_offset = timedelta(hours=tz_offset_hr) - self.dateOnly_strformat = '%d-%m' - self.dateWithYear_strformat = '%y%m%d' - self.datetime_strformat = '%Y-%m-%dT%H-%M-%S-%f' - self.register_now() - logger.info('Clock started at {}'.format(self.get_now_SGT_str())) - - def register_now(self): - self.now = datetime.now() - - def get_now_SGT(self): - ''' - returns datetime object - ''' - return self.now + self.tz_offset - - def get_now_SGT_str(self): - ''' - return date and time string - ''' - return self.get_now_SGT().strftime(self.datetime_strformat) - - def get_now_SGT_date_str(self): - ''' - returns date (MM-DD) string only - ''' - return self.get_now_SGT().strftime(self.dateOnly_strformat) - - def get_now_SGT_date_withyear_str(self): - ''' - returns date (YYMMDD) string only - ''' - return self.get_now_SGT().strftime(self.dateWithYear_strformat) diff --git a/deep_sort_realtime/utils/nms.py b/deep_sort_realtime/utils/nms.py index ca443e2..6bd98b9 100644 --- a/deep_sort_realtime/utils/nms.py +++ b/deep_sort_realtime/utils/nms.py @@ -1,5 +1,6 @@ import numpy as np + def non_max_suppression(boxes, max_bbox_overlap, scores=None): """Suppress overlapping detections. Original code from [1]_ has been adapted to include confidence score. @@ -57,7 +58,7 @@ def non_max_suppression(boxes, max_bbox_overlap, scores=None): overlap = (w * h) / area[idxs[:last]] idxs = np.delete( - idxs, np.concatenate( - ([last], np.where(overlap > max_bbox_overlap)[0]))) + idxs, np.concatenate(([last], np.where(overlap > max_bbox_overlap)[0])) + ) - return pick \ No newline at end of file + return pick diff --git a/setup.py b/setup.py index d50d5bc..22b33ba 100644 --- a/setup.py +++ b/setup.py @@ -4,25 +4,23 @@ long_description = fh.read() setup( - name='deep-sort-realtime', - version='1.0', - author="levan92", - author_email="lingevan0208@gmail.com", - description="A more realtime adaptation of Deep SORT", - long_description=long_description, - long_description_content_type="text/markdown", - url='https://github.com/levan92/deep_sort_realtime', - classifiers=[ - "Programming Language :: Python :: 3", - "License :: OSI Approved :: MIT License", - "Operating System :: OS Independent", - ], - packages=find_packages(exclude=("test",)), - package_data={ - 'deep_sort_realtime.embedder': ['weights/*'] - }, - install_requires=[ - 'numpy', - 'scipy', - ] - ) + name="deep-sort-realtime", + version="1.1", + author="levan92", + author_email="lingevan0208@gmail.com", + description="A more realtime adaptation of Deep SORT", + long_description=long_description, + long_description_content_type="text/markdown", + url="https://github.com/levan92/deep_sort_realtime", + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + ], + packages=find_packages(exclude=("test",)), + package_data={"deep_sort_realtime.embedder": ["weights/*"]}, + install_requires=[ + "numpy", + "scipy", + ], +) diff --git a/test/test_bb.py b/test/test_bb.py index a096f58..7b71e4f 100644 --- a/test/test_bb.py +++ b/test/test_bb.py @@ -1,178 +1,207 @@ import unittest import time +from datetime import datetime + class TestModule(unittest.TestCase): def test_hbb(self): from deep_sort_realtime.deepsort_tracker import DeepSort - from deep_sort_realtime.utils.clock import Clock import cv2 import numpy as np - clock = Clock() - # tracker = DeepSort(max_age = 30, nn_budget=100, nms_max_overlap=1.0, clock=clock, embedder=False) - tracker = DeepSort(max_age = 30, nn_budget=100, nms_max_overlap=1.0, clock=clock, embedder=True) + today = datetime.now().date() + tracker = DeepSort( + max_age=30, nn_budget=100, nms_max_overlap=1.0, embedder=True, today=today + ) tic = time.perf_counter() print() - print('FRAME1') - frame1 = np.ones((1080,1920,3), dtype=np.uint8) * 255 - detections1 = [ ( [0,0,50,50], 0.5, 'person' ), ([50,50, 50, 50], 0.5, 'person') ] - embeds1 = [ np.array([0.1,0.1,0.1,0.1]), np.array([-1.0,1.0,0.5,-0.5]) ] + print("FRAME1") + frame1 = np.ones((1080, 1920, 3), dtype=np.uint8) * 255 + detections1 = [ + ([0, 0, 50, 50], 0.5, "person"), + ([50, 50, 50, 50], 0.5, "person"), + ] + embeds1 = [np.array([0.1, 0.1, 0.1, 0.1]), np.array([-1.0, 1.0, 0.5, -0.5])] # tracks = tracker.update_tracks(detections1, embeds=embeds1) - tracks = tracker.update_tracks(detections1, frame=frame1) + + tracks = tracker.update_tracks( + detections1, frame=frame1, today=datetime.now().date() + ) for track in tracks: print(track.track_id) print(track.to_tlwh()) print() - print('FRAME2') + print("FRAME2") # assume new frame frame2 = frame1 - detections2 = [ ( [10,10,60,60], 0.8, 'person' ), ([60,50, 50, 50], 0.7, 'person') ] - embeds2 = [ np.array([0.1,0.1,0.1,0.1]), np.array([-1.1,1.0,0.5,-0.5]) ] + detections2 = [ + ([10, 10, 60, 60], 0.8, "person"), + ([60, 50, 50, 50], 0.7, "person"), + ] + embeds2 = [np.array([0.1, 0.1, 0.1, 0.1]), np.array([-1.1, 1.0, 0.5, -0.5])] # tracks = tracker.update_tracks(detections2, embeds=embeds2) - tracks = tracker.update_tracks(detections2, frame=frame2) + tracks = tracker.update_tracks( + detections2, frame=frame2, today=datetime.now().date() + ) for track in tracks: print(track.track_id) print(track.to_tlwh()) print() - print('FRAME3') + print("FRAME3") # assume new frame frame3 = frame1 - detections3 = [ ( [20,20,70,70], 0.8, 'person' ), ([70,50, 50, 50], 0.7, 'person') ] - embeds3 = [ np.array([0.1,0.1,0.1,0.1]), np.array([-1.1,1.0,0.5,-0.5]) ] + detections3 = [ + ([20, 20, 70, 70], 0.8, "person"), + ([70, 50, 50, 50], 0.7, "person"), + ] + embeds3 = [np.array([0.1, 0.1, 0.1, 0.1]), np.array([-1.1, 1.0, 0.5, -0.5])] # tracks = tracker.update_tracks(detections3, embeds=embeds3) - tracks = tracker.update_tracks(detections3, frame=frame3) + tracks = tracker.update_tracks( + detections3, frame=frame3, today=datetime.now().date() + ) for track in tracks: print(track.track_id) print(track.to_tlwh()) print() - print('FRAME4') + print("FRAME4") # assume new frame frame4 = frame1 - detections4 = [ ( [10,10,60,60], 0.8, 'person' )] - embeds4 = [ np.array([0.1,0.1,0.1,0.1]) ] + detections4 = [([10, 10, 60, 60], 0.8, "person")] + embeds4 = [np.array([0.1, 0.1, 0.1, 0.1])] # tracks = tracker.update_tracks(detections4, embeds=embeds4) tracks = tracker.update_tracks(detections4, frame=frame4) for track in tracks: print(track.track_id) print(track.to_tlwh()) - + toc = time.perf_counter() - print(f'Avrg Duration per update: {(toc-tic)/4}') + print(f"Avrg Duration per update: {(toc-tic)/4}") return True def test_obb(self): from deep_sort_realtime.deepsort_tracker import DeepSort - from deep_sort_realtime.utils.clock import Clock import cv2 import numpy as np - clock = Clock() - tracker = DeepSort(max_age = 30, nn_budget=100, nms_max_overlap=1.0, clock=clock, embedder=True, polygon=True) + tracker = DeepSort( + max_age=30, nn_budget=100, nms_max_overlap=1.0, embedder=True, polygon=True + ) tic = time.perf_counter() print() - print('FRAME1') - frame1 = np.ones((1080,1920,3), dtype=np.uint8) * 255 + print("FRAME1") + frame1 = np.ones((1080, 1920, 3), dtype=np.uint8) * 255 detections1 = [ - [[0,0,10,0,10,10,0,10],[20,20,30,20,30,30,20,30]], - [0,1], - [0.5,0.5] + [[0, 0, 10, 0, 10, 10, 0, 10], [20, 20, 30, 20, 30, 30, 20, 30]], + [0, 1], + [0.5, 0.5], ] tracks = tracker.update_tracks(detections1, frame=frame1) - - correct_ans = [ np.array([ 0., 0., 11., 11.]), - np.array([ 20., 20., 11., 11.]) ] + + correct_ans = [ + np.array([0.0, 0.0, 11.0, 11.0]), + np.array([20.0, 20.0, 11.0, 11.0]), + ] for track, ans in zip(tracks, correct_ans): print(track.track_id) - ltwh = track.to_ltwh() + ltwh = track.to_ltwh() print(ltwh) np.testing.assert_allclose(ltwh, ans) print() - print('FRAME2') + print("FRAME2") # assume new frame frame2 = frame1 detections2 = [ - [[0,0,10,0,15,10,0,15],[25,20,30,20,30,30,25,30]], - [0,1], - [0.5,0.6] + [[0, 0, 10, 0, 15, 10, 0, 15], [25, 20, 30, 20, 30, 30, 25, 30]], + [0, 1], + [0.5, 0.6], ] tracks = tracker.update_tracks(detections2, frame=frame2) - - correct_ans = [ np.array([ 0., 0., 15.33884298, 15.33884298]), - np.array([ 22.21844112, 20., 10.90196074, 11.]) ] + + correct_ans = [ + np.array([0.0, 0.0, 15.33884298, 15.33884298]), + np.array([22.21844112, 20.0, 10.90196074, 11.0]), + ] for track, ans in zip(tracks, correct_ans): print(track.track_id) - ltwh = track.to_ltwh() + ltwh = track.to_ltwh() print(ltwh) np.testing.assert_allclose(ltwh, ans) print() - print('FRAME3') + print("FRAME3") # assume new frame frame3 = frame1 detections3 = [ - [[0,0,10,0,15,10,10,15],[20,20,30,20,30,30,25,30]], - [0,3], - [0.5,0.6] + [[0, 0, 10, 0, 15, 10, 10, 15], [20, 20, 30, 20, 30, 30, 25, 30]], + [0, 3], + [0.5, 0.6], ] tracks = tracker.update_tracks(detections3, frame=frame3) - - correct_ans = [ np.array([ 0., 0., 16.12303476, 16.12303476]), - np.array([ 20.63971341, 20., 10.90477995, 11.]) ] + + correct_ans = [ + np.array([0.0, 0.0, 16.12303476, 16.12303476]), + np.array([20.63971341, 20.0, 10.90477995, 11.0]), + ] for track, ans in zip(tracks, correct_ans): print(track.track_id) - ltwh = track.to_ltwh() + ltwh = track.to_ltwh() print(ltwh) np.testing.assert_allclose(ltwh, ans) print() - print('FRAME4') + print("FRAME4") # assume new frame frame4 = frame1 detections4 = [ - [[0.0,5.0,15.0,5.0,15.0,10.0,10.0,25.0],[20.0,20.0,30.0,20.0,30.0,30.0,25.0,30.0]], - [3,3], - [0.9,0.6] + [ + [0.0, 5.0, 15.0, 5.0, 15.0, 10.0, 10.0, 25.0], + [20.0, 20.0, 30.0, 20.0, 30.0, 30.0, 25.0, 30.0], + ], + [3, 3], + [0.9, 0.6], ] tracks = tracker.update_tracks(detections4, frame=frame4) - - correct_ltwh_ans = [ np.array([-1.65656289, 3.48914218, 19.63792898, 19.81394538]), - np.array([ 20.10337142, 20., 10.90833262,11. ]) ] - correct_orig_ltwh_ans = [ - [0., 5., 16., 21.], - [20., 20., 11., 11.] + + correct_ltwh_ans = [ + np.array([-1.65656289, 3.48914218, 19.63792898, 19.81394538]), + np.array([20.10337142, 20.0, 10.90833262, 11.0]), ] + correct_orig_ltwh_ans = [[0.0, 5.0, 16.0, 21.0], [20.0, 20.0, 11.0, 11.0]] correct_poly_ans = detections4[0] - for track, ltwh_ans, orig_ltwh_ans, poly_ans in zip(tracks, correct_ltwh_ans, correct_orig_ltwh_ans, correct_poly_ans): + for track, ltwh_ans, orig_ltwh_ans, poly_ans in zip( + tracks, correct_ltwh_ans, correct_orig_ltwh_ans, correct_poly_ans + ): print(track.track_id) - - ltwh = track.to_ltwh() + + ltwh = track.to_ltwh() print(ltwh) np.testing.assert_allclose(ltwh, ltwh_ans) - - orig_ltwh = track.to_ltwh(orig=True) + + orig_ltwh = track.to_ltwh(orig=True) print(orig_ltwh) np.testing.assert_allclose(orig_ltwh, orig_ltwh_ans) - poly = track.get_det_supplementary() + poly = track.get_det_supplementary() print(poly) np.testing.assert_allclose(poly, poly_ans) - + toc = time.perf_counter() - print(f'Avrg Duration per update: {(toc-tic)/4}') + print(f"Avrg Duration per update: {(toc-tic)/4}") return True -if __name__ == '__main__': - unittest.main() \ No newline at end of file + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_embedder.py b/test/test_embedder.py index 1c6adb9..562fd98 100644 --- a/test/test_embedder.py +++ b/test/test_embedder.py @@ -3,40 +3,51 @@ pardir = Path(__file__).parent + def test_embedder_generic(model, thresh=10): import cv2 import numpy as np - imgpath = pardir / 'smallapple.jpg' - imgpath2 = pardir / 'rock.jpg' + imgpath = pardir / "smallapple.jpg" + imgpath2 = pardir / "rock.jpg" img = cv2.imread(str(imgpath)) small_angle = 1 image_center = tuple(np.array(img.shape[1::-1]) / 2) rot_mat = cv2.getRotationMatrix2D(image_center, small_angle, 1.0) - img_rot = cv2.warpAffine(img, rot_mat, img.shape[1::-1], flags=cv2.INTER_LINEAR, borderValue=(255,255,255)) + img_rot = cv2.warpAffine( + img, + rot_mat, + img.shape[1::-1], + flags=cv2.INTER_LINEAR, + borderValue=(255, 255, 255), + ) img2 = cv2.imread(str(imgpath2)) emb = model(max_batch_size=4) - a,b,c = emb.predict([img, img_rot, img2]) + a, b, c = emb.predict([img, img_rot, img2]) - small = np.linalg.norm(a-b) - large = np.linalg.norm(a-c) + small = np.linalg.norm(a - b) + large = np.linalg.norm(a - c) - print(f'close: {small} vs diff: {large}') - assert small < thresh,f'Small: {small} not small enough' - assert large > thresh,f'Large: {large} not large enough' + print(f"close: {small} vs diff: {large}") + assert small < thresh, f"Small: {small} not small enough" + assert large > thresh, f"Large: {large} not large enough" return True + class TestModule(unittest.TestCase): def test_embedder_torch(self): from deep_sort_realtime.embedder.embedder_pytorch import MobileNetv2_Embedder + return test_embedder_generic(MobileNetv2_Embedder) def test_embedder_tf(self): from deep_sort_realtime.embedder.embedder_tf import MobileNetv2_Embedder + return test_embedder_generic(MobileNetv2_Embedder) -if __name__ == '__main__': - unittest.main() \ No newline at end of file + +if __name__ == "__main__": + unittest.main()