base_sde_tracker.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  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. from paddlex.ppdet.core.workspace import register, serializable
  18. __all__ = ['TrackState', 'Track']
  19. class TrackState(object):
  20. """
  21. Enumeration type for the single target track state. Newly created tracks are
  22. classified as `tentative` until enough evidence has been collected. Then,
  23. the track state is changed to `confirmed`. Tracks that are no longer alive
  24. are classified as `deleted` to mark them for removal from the set of active
  25. tracks.
  26. """
  27. Tentative = 1
  28. Confirmed = 2
  29. Deleted = 3
  30. @register
  31. @serializable
  32. class Track(object):
  33. """
  34. A single target track with state space `(x, y, a, h)` and associated
  35. velocities, where `(x, y)` is the center of the bounding box, `a` is the
  36. aspect ratio and `h` is the height.
  37. Args:
  38. mean (ndarray): Mean vector of the initial state distribution.
  39. covariance (ndarray): Covariance matrix of the initial state distribution.
  40. track_id (int): A unique track identifier.
  41. n_init (int): Number of consecutive detections before the track is confirmed.
  42. The track state is set to `Deleted` if a miss occurs within the first
  43. `n_init` frames.
  44. max_age (int): The maximum number of consecutive misses before the track
  45. state is set to `Deleted`.
  46. feature (Optional[ndarray]): Feature vector of the detection this track
  47. originates from. If not None, this feature is added to the `features` cache.
  48. Attributes:
  49. hits (int): Total number of measurement updates.
  50. age (int): Total number of frames since first occurance.
  51. time_since_update (int): Total number of frames since last measurement
  52. update.
  53. state (TrackState): The current track state.
  54. features (List[ndarray]): A cache of features. On each measurement update,
  55. the associated feature vector is added to this list.
  56. """
  57. def __init__(self,
  58. mean,
  59. covariance,
  60. track_id,
  61. n_init,
  62. max_age,
  63. feature=None):
  64. self.mean = mean
  65. self.covariance = covariance
  66. self.track_id = track_id
  67. self.hits = 1
  68. self.age = 1
  69. self.time_since_update = 0
  70. self.state = TrackState.Tentative
  71. self.features = []
  72. if feature is not None:
  73. self.features.append(feature)
  74. self._n_init = n_init
  75. self._max_age = max_age
  76. def to_tlwh(self):
  77. """Get position in format `(top left x, top left y, width, height)`."""
  78. ret = self.mean[:4].copy()
  79. ret[2] *= ret[3]
  80. ret[:2] -= ret[2:] / 2
  81. return ret
  82. def to_tlbr(self):
  83. """Get position in bounding box format `(min x, miny, max x, max y)`."""
  84. ret = self.to_tlwh()
  85. ret[2:] = ret[:2] + ret[2:]
  86. return ret
  87. def predict(self, kalman_filter):
  88. """
  89. Propagate the state distribution to the current time step using a Kalman
  90. filter prediction step.
  91. """
  92. self.mean, self.covariance = kalman_filter.predict(self.mean,
  93. self.covariance)
  94. self.age += 1
  95. self.time_since_update += 1
  96. def update(self, kalman_filter, detection):
  97. """
  98. Perform Kalman filter measurement update step and update the associated
  99. detection feature cache.
  100. """
  101. self.mean, self.covariance = kalman_filter.update(self.mean,
  102. self.covariance,
  103. detection.to_xyah())
  104. self.features.append(detection.feature)
  105. self.hits += 1
  106. self.time_since_update = 0
  107. if self.state == TrackState.Tentative and self.hits >= self._n_init:
  108. self.state = TrackState.Confirmed
  109. def mark_missed(self):
  110. """Mark this track as missed (no association at the current time step).
  111. """
  112. if self.state == TrackState.Tentative:
  113. self.state = TrackState.Deleted
  114. elif self.time_since_update > self._max_age:
  115. self.state = TrackState.Deleted
  116. def is_tentative(self):
  117. """Returns True if this track is tentative (unconfirmed)."""
  118. return self.state == TrackState.Tentative
  119. def is_confirmed(self):
  120. """Returns True if this track is confirmed."""
  121. return self.state == TrackState.Confirmed
  122. def is_deleted(self):
  123. """Returns True if this track is dead and should be deleted."""
  124. return self.state == TrackState.Deleted