1
+ """
2
+ Hybrid Sort
3
+ """
4
+
5
+ import numpy as np
6
+ from collections import deque
7
+ from .basetrack import BaseTrack , TrackState
8
+ from .tracklet import Tracklet , Tracklet_w_velocity_four_corner
9
+ from .matching import *
10
+
11
+ from cython_bbox import bbox_overlaps as bbox_ious
12
+
13
+ class HybridSortTracker (object ):
14
+ def __init__ (self , args , frame_rate = 30 ):
15
+ self .tracked_tracklets = [] # type: list[Tracklet]
16
+ self .lost_tracklets = [] # type: list[Tracklet]
17
+ self .removed_tracklets = [] # type: list[Tracklet]
18
+
19
+ self .frame_id = 0
20
+ self .args = args
21
+
22
+ self .det_thresh = args .conf_thresh + 0.1
23
+ self .buffer_size = int (frame_rate / 30.0 * args .track_buffer )
24
+ self .max_time_lost = self .buffer_size
25
+
26
+ self .motion = args .kalman_format
27
+
28
+ self .delta_t = 3
29
+
30
+ @staticmethod
31
+ def k_previous_obs (observations , cur_age , k ):
32
+ if len (observations ) == 0 :
33
+ return [- 1 , - 1 , - 1 , - 1 , - 1 ]
34
+ for i in range (k ):
35
+ dt = k - i
36
+ if cur_age - dt in observations :
37
+ return observations [cur_age - dt ]
38
+ max_age = max (observations .keys ())
39
+ return observations [max_age ]
40
+
41
+ def update (self , output_results , img , ori_img ):
42
+ """
43
+ output_results: processed detections (scale to original size) tlbr format
44
+ """
45
+
46
+ self .frame_id += 1
47
+ activated_tracklets = []
48
+ refind_tracklets = []
49
+ lost_tracklets = []
50
+ removed_tracklets = []
51
+
52
+ scores = output_results [:, 4 ]
53
+ bboxes = output_results [:, :4 ]
54
+ categories = output_results [:, - 1 ]
55
+
56
+ remain_inds = scores > self .args .conf_thresh
57
+ inds_low = scores > 0.1
58
+ inds_high = scores < self .args .conf_thresh
59
+
60
+ inds_second = np .logical_and (inds_low , inds_high )
61
+ dets_second = bboxes [inds_second ]
62
+ dets = bboxes [remain_inds ]
63
+
64
+ cates = categories [remain_inds ]
65
+ cates_second = categories [inds_second ]
66
+
67
+ scores_keep = scores [remain_inds ]
68
+ scores_second = scores [inds_second ]
69
+
70
+ if len (dets ) > 0 :
71
+ '''Detections'''
72
+ detections = [Tracklet_w_velocity_four_corner (tlwh , s , cate , motion = self .motion ) for
73
+ (tlwh , s , cate ) in zip (dets , scores_keep , cates )]
74
+ else :
75
+ detections = []
76
+
77
+ ''' Add newly detected tracklets to tracked_tracklets'''
78
+ unconfirmed = []
79
+ tracked_tracklets = [] # type: list[Tracklet]
80
+ for track in self .tracked_tracklets :
81
+ if not track .is_activated :
82
+ unconfirmed .append (track )
83
+ else :
84
+ tracked_tracklets .append (track )
85
+
86
+ ''' Step 2: First association, Weak Cues (four corner confidence and score)'''
87
+ tracklet_pool = joint_tracklets (tracked_tracklets , self .lost_tracklets )
88
+
89
+ velocities = np .array (
90
+ [trk .get_velocity () for trk in tracklet_pool ]) # (N, 4, 2)
91
+
92
+ # last observation, obervation-centric
93
+ # last_boxes = np.array([trk.last_observation for trk in tracklet_pool])
94
+
95
+ # historical observations
96
+ k_observations = np .array (
97
+ [self .k_previous_obs (trk .observations , trk .age , self .delta_t ) for trk in tracklet_pool ])
98
+
99
+
100
+ # Predict the current location with Kalman
101
+ for tracklet in tracklet_pool :
102
+ tracklet .predict ()
103
+
104
+ # weak cues cost matrix (hmiou + four corner velocity) and assignment
105
+ matches , u_track , u_detection = association_weak_cues (
106
+ tracklets = tracklet_pool , detections = detections , velocities = velocities ,
107
+ previous_obs = k_observations , vdc_weight = 0.05 )
108
+
109
+ for itracked , idet in matches :
110
+ track = tracklet_pool [itracked ]
111
+ det = detections [idet ]
112
+ if track .state == TrackState .Tracked :
113
+ track .update (detections [idet ], self .frame_id )
114
+ activated_tracklets .append (track )
115
+ else :
116
+ track .re_activate (det , self .frame_id , new_id = False )
117
+ refind_tracklets .append (track )
118
+
119
+ ''' Step 3: Second association, with low score detection boxes'''
120
+ # association the untrack to the low score detections
121
+ if len (dets_second ) > 0 :
122
+ '''Detections'''
123
+ detections_second = [Tracklet_w_velocity_four_corner (tlwh , s , cate , motion = self .motion ) for
124
+ (tlwh , s , cate ) in zip (dets_second , scores_second , cates_second )]
125
+ else :
126
+ detections_second = []
127
+ r_tracked_tracklets = [tracklet_pool [i ] for i in u_track if tracklet_pool [i ].state == TrackState .Tracked ]
128
+
129
+ dists = hm_iou_distance (r_tracked_tracklets , detections_second ) - score_distance (r_tracked_tracklets , detections_second )
130
+
131
+ matches , u_track , u_detection_second = linear_assignment (dists , thresh = 0.5 )
132
+ for itracked , idet in matches :
133
+ track = r_tracked_tracklets [itracked ]
134
+ det = detections_second [idet ]
135
+ if track .state == TrackState .Tracked :
136
+ track .update (det , self .frame_id )
137
+ activated_tracklets .append (track )
138
+ else :
139
+ track .re_activate (det , self .frame_id , new_id = False )
140
+ refind_tracklets .append (track )
141
+
142
+
143
+ '''Step 4: Third association, match high-conf remain detections with last observation of tracks'''
144
+ r_tracked_tracklets = [r_tracked_tracklets [i ] for i in u_track ] # remain tracklets from last step
145
+ r_detections = [detections [i ] for i in u_detection ] # high-conf remain detections
146
+
147
+ dists = hm_iou_distance (atracks = [t .last_observation [: 4 ] for t in r_tracked_tracklets ], # parse bbox directly
148
+ btracks = [d .tlbr for d in r_detections ])
149
+
150
+ matches , u_track , u_detection = linear_assignment (dists , thresh = 0.5 )
151
+
152
+ for itracked , idet in matches :
153
+ track = r_tracked_tracklets [itracked ]
154
+ det = r_detections [idet ]
155
+ if track .state == TrackState .Tracked :
156
+ track .update (det , self .frame_id )
157
+ activated_tracklets .append (track )
158
+ else :
159
+ track .re_activate (det , self .frame_id , new_id = False )
160
+ refind_tracklets .append (track )
161
+
162
+ # for tracks still failed, mark lost
163
+ for it in u_track :
164
+ track = r_tracked_tracklets [it ]
165
+ if not track .state == TrackState .Lost :
166
+ track .mark_lost ()
167
+ lost_tracklets .append (track )
168
+
169
+
170
+ '''Deal with unconfirmed tracks, usually tracks with only one beginning frame'''
171
+ detections = [r_detections [i ] for i in u_detection ]
172
+ dists = iou_distance (unconfirmed , detections )
173
+
174
+ matches , u_unconfirmed , u_detection = linear_assignment (dists , thresh = 0.7 )
175
+
176
+ for itracked , idet in matches :
177
+ unconfirmed [itracked ].update (detections [idet ], self .frame_id )
178
+ activated_tracklets .append (unconfirmed [itracked ])
179
+ for it in u_unconfirmed :
180
+ track = unconfirmed [it ]
181
+ track .mark_removed ()
182
+ removed_tracklets .append (track )
183
+
184
+ """ Step 4: Init new tracklets"""
185
+ for inew in u_detection :
186
+ track = detections [inew ]
187
+ if track .score < self .det_thresh :
188
+ continue
189
+ track .activate (self .frame_id )
190
+ activated_tracklets .append (track )
191
+
192
+ """ Step 5: Update state"""
193
+ for track in self .lost_tracklets :
194
+ if self .frame_id - track .end_frame > self .max_time_lost :
195
+ track .mark_removed ()
196
+ removed_tracklets .append (track )
197
+
198
+ self .tracked_tracklets = [t for t in self .tracked_tracklets if t .state == TrackState .Tracked ]
199
+ self .tracked_tracklets = joint_tracklets (self .tracked_tracklets , activated_tracklets )
200
+ self .tracked_tracklets = joint_tracklets (self .tracked_tracklets , refind_tracklets )
201
+ self .lost_tracklets = sub_tracklets (self .lost_tracklets , self .tracked_tracklets )
202
+ self .lost_tracklets .extend (lost_tracklets )
203
+ self .lost_tracklets = sub_tracklets (self .lost_tracklets , self .removed_tracklets )
204
+ self .removed_tracklets .extend (removed_tracklets )
205
+ self .tracked_tracklets , self .lost_tracklets = remove_duplicate_tracklets (self .tracked_tracklets , self .lost_tracklets )
206
+ # get scores of lost tracks
207
+ output_tracklets = [track for track in self .tracked_tracklets if track .is_activated ]
208
+
209
+ return output_tracklets
210
+
211
+
212
+
213
+
214
+ def joint_tracklets (tlista , tlistb ):
215
+ exists = {}
216
+ res = []
217
+ for t in tlista :
218
+ exists [t .track_id ] = 1
219
+ res .append (t )
220
+ for t in tlistb :
221
+ tid = t .track_id
222
+ if not exists .get (tid , 0 ):
223
+ exists [tid ] = 1
224
+ res .append (t )
225
+ return res
226
+
227
+
228
+ def sub_tracklets (tlista , tlistb ):
229
+ tracklets = {}
230
+ for t in tlista :
231
+ tracklets [t .track_id ] = t
232
+ for t in tlistb :
233
+ tid = t .track_id
234
+ if tracklets .get (tid , 0 ):
235
+ del tracklets [tid ]
236
+ return list (tracklets .values ())
237
+
238
+
239
+ def remove_duplicate_tracklets (trackletsa , trackletsb ):
240
+ pdist = iou_distance (trackletsa , trackletsb )
241
+ pairs = np .where (pdist < 0.15 )
242
+ dupa , dupb = list (), list ()
243
+ for p , q in zip (* pairs ):
244
+ timep = trackletsa [p ].frame_id - trackletsa [p ].start_frame
245
+ timeq = trackletsb [q ].frame_id - trackletsb [q ].start_frame
246
+ if timep > timeq :
247
+ dupb .append (q )
248
+ else :
249
+ dupa .append (p )
250
+ resa = [t for i , t in enumerate (trackletsa ) if not i in dupa ]
251
+ resb = [t for i , t in enumerate (trackletsb ) if not i in dupb ]
252
+ return resa , resb
0 commit comments