base_sde_tracker.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """
  15. This code is borrow from https://github.com/nwojke/deep_sort/blob/master/deep_sort/track.py
  16. """
  17. import numpy as np
  18. from paddlex.ppdet.core.workspace import register, serializable
  19. __all__ = ['TrackState', 'Track']
  20. class TrackState(object):
  21. """
  22. Enumeration type for the single target track state. Newly created tracks are
  23. classified as `tentative` until enough evidence has been collected. Then,
  24. the track state is changed to `confirmed`. Tracks that are no longer alive
  25. are classified as `deleted` to mark them for removal from the set of active
  26. tracks.
  27. """
  28. Tentative = 1
  29. Confirmed = 2
  30. Deleted = 3
  31. @register
  32. @serializable
  33. class Track(object):
  34. """
  35. A single target track with state space `(x, y, a, h)` and associated
  36. velocities, where `(x, y)` is the center of the bounding box, `a` is the
  37. aspect ratio and `h` is the height.
  38. Args:
  39. mean (ndarray): Mean vector of the initial state distribution.
  40. covariance (ndarray): Covariance matrix of the initial state distribution.
  41. track_id (int): A unique track identifier.
  42. n_init (int): Number of consecutive detections before the track is confirmed.
  43. The track state is set to `Deleted` if a miss occurs within the first
  44. `n_init` frames.
  45. max_age (int): The maximum number of consecutive misses before the track
  46. state is set to `Deleted`.
  47. feature (Optional[ndarray]): Feature vector of the detection this track
  48. originates from. If not None, this feature is added to the `features` cache.
  49. Attributes:
  50. hits (int): Total number of measurement updates.
  51. age (int): Total number of frames since first occurance.
  52. time_since_update (int): Total number of frames since last measurement
  53. update.
  54. state (TrackState): The current track state.
  55. features (List[ndarray]): A cache of features. On each measurement update,
  56. the associated feature vector is added to this list.
  57. """
  58. def __init__(self,
  59. mean,
  60. covariance,
  61. track_id,
  62. n_init,
  63. max_age,
  64. feature=None):
  65. self.mean = mean
  66. self.covariance = covariance
  67. self.track_id = track_id
  68. self.hits = 1
  69. self.age = 1
  70. self.time_since_update = 0
  71. self.state = TrackState.Tentative
  72. self.features = []
  73. if feature is not None:
  74. self.features.append(feature)
  75. self._n_init = n_init
  76. self._max_age = max_age
  77. def to_tlwh(self):
  78. """Get position in format `(top left x, top left y, width, height)`."""
  79. ret = self.mean[:4].copy()
  80. ret[2] *= ret[3]
  81. ret[:2] -= ret[2:] / 2
  82. return ret
  83. def to_tlbr(self):
  84. """Get position in bounding box format `(min x, miny, max x, max y)`."""
  85. ret = self.to_tlwh()
  86. ret[2:] = ret[:2] + ret[2:]
  87. return ret
  88. def predict(self, kalman_filter):
  89. """
  90. Propagate the state distribution to the current time step using a Kalman
  91. filter prediction step.
  92. """
  93. self.mean, self.covariance = kalman_filter.predict(self.mean,
  94. self.covariance)
  95. self.age += 1
  96. self.time_since_update += 1
  97. def update(self, kalman_filter, detection):
  98. """
  99. Perform Kalman filter measurement update step and update the associated
  100. detection feature cache.
  101. """
  102. self.mean, self.covariance = kalman_filter.update(self.mean,
  103. self.covariance,
  104. detection.to_xyah())
  105. self.features.append(detection.feature)
  106. self.hits += 1
  107. self.time_since_update = 0
  108. if self.state == TrackState.Tentative and self.hits >= self._n_init:
  109. self.state = TrackState.Confirmed
  110. def mark_missed(self):
  111. """Mark this track as missed (no association at the current time step).
  112. """
  113. if self.state == TrackState.Tentative:
  114. self.state = TrackState.Deleted
  115. elif self.time_since_update > self._max_age:
  116. self.state = TrackState.Deleted
  117. def is_tentative(self):
  118. """Returns True if this track is tentative (unconfirmed)."""
  119. return self.state == TrackState.Tentative
  120. def is_confirmed(self):
  121. """Returns True if this track is confirmed."""
  122. return self.state == TrackState.Confirmed
  123. def is_deleted(self):
  124. """Returns True if this track is dead and should be deleted."""
  125. return self.state == TrackState.Deleted