| 17 | |
| 18 | |
| 19 | class STrack(BaseTrack): |
| 20 | |
| 21 | def __init__(self, tlwh, score, max_n_features=100, from_det=True): |
| 22 | |
| 23 | # wait activate |
| 24 | self._tlwh = np.asarray(tlwh, dtype=np.float) |
| 25 | self.kalman_filter = None |
| 26 | self.mean, self.covariance = None, None |
| 27 | self.is_activated = False |
| 28 | |
| 29 | self.score = score |
| 30 | self.max_n_features = max_n_features |
| 31 | self.curr_feature = None |
| 32 | self.last_feature = None |
| 33 | self.features = deque([], maxlen=self.max_n_features) |
| 34 | |
| 35 | # classification |
| 36 | self.from_det = from_det |
| 37 | self.tracklet_len = 0 |
| 38 | self.time_by_tracking = 0 |
| 39 | |
| 40 | # self-tracking |
| 41 | self.tracker = None |
| 42 | |
| 43 | def set_feature(self, feature): |
| 44 | if feature is None: |
| 45 | return False |
| 46 | self.features.append(feature) |
| 47 | self.curr_feature = feature |
| 48 | self.last_feature = feature |
| 49 | # self._p_feature = 0 |
| 50 | return True |
| 51 | |
| 52 | def predict(self): |
| 53 | if self.time_since_update > 0: |
| 54 | self.tracklet_len = 0 |
| 55 | |
| 56 | self.time_since_update += 1 |
| 57 | |
| 58 | mean_state = self.mean.copy() |
| 59 | if self.state != TrackState.Tracked: |
| 60 | mean_state[7] = 0 |
| 61 | self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance) |
| 62 | |
| 63 | if self.tracker: |
| 64 | self.tracker.update_roi(self.tlwh) |
| 65 | |
| 66 | def self_tracking(self, image): |
| 67 | tlwh = self.tracker.predict(image) if self.tracker else self.tlwh |
| 68 | return tlwh |
| 69 | |
| 70 | def activate(self, kalman_filter, frame_id, image): |
| 71 | """Start a new tracklet""" |
| 72 | self.kalman_filter = kalman_filter # type: KalmanFilter |
| 73 | self.track_id = self.next_id() |
| 74 | # cx, cy, aspect_ratio, height, dx, dy, da, dh |
| 75 | self.mean, self.covariance = self.kalman_filter.initiate(self.tlwh_to_xyah(self._tlwh)) |
| 76 | |