1+ import numpy as np
2+ from scipy .optimize import linear_sum_assignment
3+ from filterpy .kalman import KalmanFilter
4+
5+ class KalmanBoxTracker (object ):
6+ count = 0
7+
8+ def __init__ (self , bbox ):
9+ self .kf = KalmanFilter (dim_x = 7 , dim_z = 4 )
10+ self .kf .F = np .array ([[1 , 0 , 0 , 0 , 1 , 0 , 0 ], [0 , 1 , 0 , 0 , 0 , 1 , 0 ], [0 , 0 , 1 , 0 , 0 , 0 , 1 ], [0 , 0 , 0 , 1 , 0 , 0 , 0 ], [0 , 0 , 0 , 0 , 1 , 0 , 0 ], [0 , 0 , 0 , 0 , 0 , 1 , 0 ], [0 , 0 , 0 , 0 , 0 , 0 , 1 ]])
11+ self .kf .H = np .array ([[1 , 0 , 0 , 0 , 0 , 0 , 0 ], [0 , 1 , 0 , 0 , 0 , 0 , 0 ], [0 , 0 , 1 , 0 , 0 , 0 , 0 ], [0 , 0 , 0 , 1 , 0 , 0 , 0 ]])
12+ self .kf .R [2 :, 2 :] *= 10.0
13+ self .kf .P [4 :, 4 :] *= 1000.0
14+ self .kf .P *= 10.0
15+ self .kf .Q [- 1 , - 1 ] *= 0.01
16+ self .kf .Q [4 :, 4 :] *= 0.01
17+ self .kf .x [:4 ] = self .convert_bbox_to_z (bbox )
18+ self .time_since_update = 0
19+ self .id = KalmanBoxTracker .count
20+ KalmanBoxTracker .count += 1
21+ self .history = []
22+ self .hits = 0
23+ self .hit_streak = 0
24+ self .age = 0
25+
26+ def update (self , bbox ):
27+ self .time_since_update = 0
28+ self .history = []
29+ self .hits += 1
30+ self .hit_streak += 1
31+ self .kf .update (self .convert_bbox_to_z (bbox ))
32+
33+ def predict (self ):
34+ if self .kf .x [6 ] + self .kf .x [2 ] <= 0 :
35+ self .kf .x [6 ] *= 0.0
36+ self .kf .predict ()
37+ self .age += 1
38+ if self .time_since_update > 0 :
39+ self .hit_streak = 0
40+ self .time_since_update += 1
41+ self .history .append (self .convert_x_to_bbox (self .kf .x ))
42+ return self .history [- 1 ]
43+
44+ def get_state (self ):
45+ return self .convert_x_to_bbox (self .kf .x )
46+
47+ @staticmethod
48+ def convert_bbox_to_z (bbox ):
49+ w = bbox [2 ] - bbox [0 ]
50+ h = bbox [3 ] - bbox [1 ]
51+ x = bbox [0 ] + w / 2.0
52+ y = bbox [1 ] + h / 2.0
53+ s = w * h
54+ r = w / float (h )
55+ return np .array ([x , y , s , r ]).reshape ((4 , 1 ))
56+
57+ @staticmethod
58+ def convert_x_to_bbox (x , score = None ):
59+ w = np .sqrt (x [2 ] * x [3 ])
60+ h = x [2 ] / w
61+ if score is None :
62+ return np .array ([x [0 ] - w / 2.0 , x [1 ] - h / 2.0 , x [0 ] + w / 2.0 , x [1 ] + h / 2.0 ]).reshape ((1 , 4 ))
63+ return np .array ([x [0 ] - w / 2.0 , x [1 ] - h / 2.0 , x [0 ] + w / 2.0 , x [1 ] + h / 2.0 , score ]).reshape ((1 , 5 ))
64+
65+ class SimpleDeepSORT :
66+
67+ def __init__ (self , max_age = 30 , min_hits = 0 , iou_threshold = 0.3 ):
68+ self .max_age = max_age
69+ self .min_hits = min_hits
70+ self .iou_threshold = iou_threshold
71+ self .trackers = []
72+ self .frame_count = 0
73+
74+ @staticmethod
75+ def convert_bbox_xyxy_to_xywh (bbox ):
76+ """
77+ Convert bounding boxes from format (x1, y1, x2, y2) to (center x, center y, width, height).
78+
79+ Args:
80+ bbox (np.array): Bounding boxes in format (x1, y1, x2, y2), can be a single box or multiple boxes.
81+
82+ Returns:
83+ np.array: Bounding boxes in format (x, y, w, h)
84+ """
85+ bbox = np .array (bbox )
86+ if bbox .ndim == 1 :
87+ bbox = bbox [np .newaxis , :]
88+ x1 , y1 , x2 , y2 = (bbox [:, 0 ], bbox [:, 1 ], bbox [:, 2 ], bbox [:, 3 ])
89+ x = (x1 + x2 ) / 2
90+ y = (y1 + y2 ) / 2
91+ w = x2 - x1
92+ h = y2 - y1
93+ return np .stack ((x , y , w , h ), axis = - 1 )
94+
95+ @staticmethod
96+ def convert_bbox_xywh_to_xyxy (bbox ):
97+ """
98+ Convert bounding boxes from format (center x, center y, width, height) to (x1, y1, x2, y2).
99+
100+ Args:
101+ bbox (np.array): Bounding boxes in format (x, y, w, h), can be a single box or multiple boxes.
102+
103+ Returns:
104+ np.array: Bounding boxes in format (x1, y1, x2, y2)
105+ """
106+ bbox = np .array (bbox )
107+ if bbox .ndim == 1 :
108+ bbox = bbox [np .newaxis , :]
109+ x , y , w , h = (bbox [:, 0 ], bbox [:, 1 ], bbox [:, 2 ], bbox [:, 3 ])
110+ x1 = x - w / 2
111+ y1 = y - h / 2
112+ x2 = x + w / 2
113+ y2 = y + h / 2
114+ return np .stack ((x1 , y1 , x2 , y2 ), axis = - 1 )
115+
116+ def update (self , dets_xywh ):
117+ dets = self .convert_bbox_xywh_to_xyxy (dets_xywh )
118+ self .frame_count += 1
119+ trks = np .zeros ((len (self .trackers ), 5 ))
120+ to_del = []
121+ ret = []
122+ for t , trk in enumerate (trks ):
123+ pos = self .trackers [t ].predict ()[0 ]
124+ trk [:] = [pos [0 ], pos [1 ], pos [2 ], pos [3 ], 0 ]
125+ if np .any (np .isnan (pos )):
126+ to_del .append (t )
127+ trks = np .ma .compress_rows (np .ma .masked_invalid (trks ))
128+ for t in reversed (to_del ):
129+ self .trackers .pop (t )
130+ matched , unmatched_dets , unmatched_trks = self .associate_detections_to_trackers (dets , trks )
131+ for m in matched :
132+ self .trackers [m [1 ]].update (dets [m [0 ], :])
133+ for i in unmatched_dets :
134+ trk = KalmanBoxTracker (dets [i , :])
135+ self .trackers .append (trk )
136+ i = len (self .trackers )
137+ for trk in reversed (self .trackers ):
138+ d = trk .history [- 1 ] if len (trk .history ) > 0 else trk .get_state ()[0 ]
139+ if trk .time_since_update < 1 and (trk .hit_streak >= self .min_hits or self .frame_count <= self .min_hits ):
140+ ret .append (np .concatenate ((d , [trk .id + 1 ])).reshape (1 , - 1 ))
141+ i -= 1
142+ if trk .time_since_update > self .max_age :
143+ self .trackers .pop (i )
144+ if len (ret ) > 0 :
145+ return (self .convert_bbox_xyxy_to_xywh (np .concatenate (ret )[:, :4 ]), np .concatenate (ret )[:, 4 ])
146+ return (np .empty ((0 , 4 )), np .empty (0 , dtype = int ))
147+
148+ def associate_detections_to_trackers (self , detections , trackers ):
149+ if len (trackers ) == 0 :
150+ return (np .empty ((0 , 2 ), dtype = int ), np .arange (len (detections )), np .empty ((0 , 5 ), dtype = int ))
151+ matched_indices = linear_sum_assignment (- np .ones ((len (detections ), len (trackers ))))
152+ matched_indices = np .asarray (matched_indices )
153+ matched_indices = np .transpose (matched_indices )
154+ unmatched_detections = []
155+ for d , det in enumerate (detections ):
156+ if d not in matched_indices [:, 0 ]:
157+ unmatched_detections .append (d )
158+ unmatched_trackers = []
159+ for t , trk in enumerate (trackers ):
160+ if t not in matched_indices [:, 1 ]:
161+ unmatched_trackers .append (t )
162+ matches = []
163+ for m in matched_indices :
164+ matches .append (m .reshape (1 , 2 ))
165+ if len (matches ) == 0 :
166+ matches = np .empty ((0 , 2 ), dtype = int )
167+ else :
168+ matches = np .concatenate (matches , axis = 0 )
169+ return (matches , np .array (unmatched_detections ), np .array (unmatched_trackers ))
170+
171+ def _iou (self , bb_test , bb_gt ):
172+ xx1 = np .maximum (bb_test [0 ], bb_gt [0 ])
173+ yy1 = np .maximum (bb_test [1 ], bb_gt [1 ])
174+ xx2 = np .minimum (bb_test [2 ], bb_gt [2 ])
175+ yy2 = np .minimum (bb_test [3 ], bb_gt [3 ])
176+ w = np .maximum (0.0 , xx2 - xx1 )
177+ h = np .maximum (0.0 , yy2 - yy1 )
178+ wh = w * h
179+ o = wh / ((bb_test [2 ] - bb_test [0 ]) * (bb_test [3 ] - bb_test [1 ]) + (bb_gt [2 ] - bb_gt [0 ]) * (bb_gt [3 ] - bb_gt [1 ]) - wh )
180+ return o
0 commit comments