99import csv
1010
1111
12- def load_mot (detections ):
12+ visdrone_classes = {'car' : 4 , 'bus' : 9 , 'truck' : 6 , 'pedestrian' : 1 , 'van' : 5 }
13+
14+
15+ def load_mot (detections , nms_overlap_thresh = None , with_classes = True , nms_per_class = False ):
1316 """
1417 Loads detections stored in a mot-challenge like formatted CSV or numpy array (fieldNames = ['frame', 'id', 'x', 'y',
1518 'w', 'h', 'score']).
1619
1720 Args:
18- detections
21+ detections (str, numpy.ndarray): path to csv file containing the detections or numpy array containing them.
22+ nms_overlap_thresh (float, optional): perform non-maximum suppression on the input detections with this thrshold.
23+ no nms is performed if this parameter is not specified.
24+ with_classes (bool, optional): indicates if the detections have classes or not. set to false for motchallange.
25+ nms_per_class (bool, optional): perform non-maximum suppression for each class separately
1926
2027 Returns:
2128 list: list containing the detections for each frame.
2229 """
30+ if nms_overlap_thresh :
31+ assert with_classes , "currently only works with classes available"
2332
2433 data = []
2534 if type (detections ) is str :
2635 raw = np .genfromtxt (detections , delimiter = ',' , dtype = np .float32 )
36+ if np .isnan (raw ).all ():
37+ raw = np .genfromtxt (detections , delimiter = ' ' , dtype = np .float32 )
38+
2739 else :
2840 # assume it is an array
2941 assert isinstance (detections , np .ndarray ), "only numpy arrays or *.csv paths are supported as detections."
@@ -34,16 +46,127 @@ def load_mot(detections):
3446 idx = raw [:, 0 ] == i
3547 bbox = raw [idx , 2 :6 ]
3648 bbox [:, 2 :4 ] += bbox [:, 0 :2 ] # x1, y1, w, h -> x1, y1, x2, y2
49+ bbox -= 1 # correct 1,1 matlab offset
3750 scores = raw [idx , 6 ]
51+
52+ if with_classes :
53+ classes = raw [idx , 7 ]
54+
55+ bbox_filtered = None
56+ scores_filtered = None
57+ classes_filtered = None
58+ for coi in visdrone_classes :
59+ cids = classes == visdrone_classes [coi ]
60+ if nms_per_class and nms_overlap_thresh :
61+ bbox_tmp , scores_tmp = nms (bbox [cids ], scores [cids ], nms_overlap_thresh )
62+ else :
63+ bbox_tmp , scores_tmp = bbox [cids ], scores [cids ]
64+
65+ if bbox_filtered is None :
66+ bbox_filtered = bbox_tmp
67+ scores_filtered = scores_tmp
68+ classes_filtered = [coi ]* bbox_filtered .shape [0 ]
69+ elif len (bbox_tmp ) > 0 :
70+ bbox_filtered = np .vstack ((bbox_filtered , bbox_tmp ))
71+ scores_filtered = np .hstack ((scores_filtered , scores_tmp ))
72+ classes_filtered += [coi ] * bbox_tmp .shape [0 ]
73+
74+ if bbox_filtered is not None :
75+ bbox = bbox_filtered
76+ scores = scores_filtered
77+ classes = classes_filtered
78+
79+ if nms_per_class is False and nms_overlap_thresh :
80+ bbox , scores , classes = nms (bbox , scores , nms_overlap_thresh , np .array (classes ))
81+
82+ else :
83+ classes = ['pedestrian' ]* bbox .shape [0 ]
84+
3885 dets = []
39- for bb , s in zip (bbox , scores ):
40- dets .append ({'bbox' : (bb [0 ], bb [1 ], bb [2 ], bb [3 ]), 'score' : s })
86+ for bb , s , c in zip (bbox , scores , classes ):
87+ dets .append ({'bbox' : (bb [0 ], bb [1 ], bb [2 ], bb [3 ]), 'score' : s , 'class' : c })
4188 data .append (dets )
4289
4390 return data
4491
4592
46- def save_to_csv (out_path , tracks ):
93+ def nms (boxes , scores , overlapThresh , classes = None ):
94+ """
95+ perform non-maximum suppression. based on Malisiewicz et al.
96+ Args:
97+ boxes (numpy.ndarray): boxes to process
98+ scores (numpy.ndarray): corresponding scores for each box
99+ overlapThresh (float): overlap threshold for boxes to merge
100+ classes (numpy.ndarray, optional): class ids for each box.
101+
102+ Returns:
103+ (tuple): tuple containing:
104+
105+ boxes (list): nms boxes
106+ scores (list): nms scores
107+ classes (list, optional): nms classes if specified
108+ """
109+ # # if there are no boxes, return an empty list
110+ # if len(boxes) == 0:
111+ # return [], [], [] if classes else [], []
112+
113+ # if the bounding boxes integers, convert them to floats --
114+ # this is important since we'll be doing a bunch of divisions
115+ if boxes .dtype .kind == "i" :
116+ boxes = boxes .astype ("float" )
117+
118+ if scores .dtype .kind == "i" :
119+ scores = scores .astype ("float" )
120+
121+ # initialize the list of picked indexes
122+ pick = []
123+
124+ # grab the coordinates of the bounding boxes
125+ x1 = boxes [:, 0 ]
126+ y1 = boxes [:, 1 ]
127+ x2 = boxes [:, 2 ]
128+ y2 = boxes [:, 3 ]
129+ #score = boxes[:, 4]
130+ # compute the area of the bounding boxes and sort the bounding
131+ # boxes by the bottom-right y-coordinate of the bounding box
132+ area = (x2 - x1 + 1 ) * (y2 - y1 + 1 )
133+ idxs = np .argsort (scores )
134+
135+ # keep looping while some indexes still remain in the indexes
136+ # list
137+ while len (idxs ) > 0 :
138+ # grab the last index in the indexes list and add the
139+ # index value to the list of picked indexes
140+ last = len (idxs ) - 1
141+ i = idxs [last ]
142+ pick .append (i )
143+
144+ # find the largest (x, y) coordinates for the start of
145+ # the bounding box and the smallest (x, y) coordinates
146+ # for the end of the bounding box
147+ xx1 = np .maximum (x1 [i ], x1 [idxs [:last ]])
148+ yy1 = np .maximum (y1 [i ], y1 [idxs [:last ]])
149+ xx2 = np .minimum (x2 [i ], x2 [idxs [:last ]])
150+ yy2 = np .minimum (y2 [i ], y2 [idxs [:last ]])
151+
152+ # compute the width and height of the bounding box
153+ w = np .maximum (0 , xx2 - xx1 + 1 )
154+ h = np .maximum (0 , yy2 - yy1 + 1 )
155+
156+ # compute the ratio of overlap
157+ overlap = (w * h ) / area [idxs [:last ]]
158+
159+ # delete all indexes from the index list that have
160+ idxs = np .delete (idxs , np .concatenate (([last ],
161+ np .where (overlap > overlapThresh )[0 ])))
162+
163+ if classes is not None :
164+ return boxes [pick ], scores [pick ], classes [pick ]
165+ else :
166+ return boxes [pick ], scores [pick ]
167+
168+
169+ def save_to_csv (out_path , tracks , fmt = 'motchallenge' ):
47170 """
48171 Saves tracks to a CSV file.
49172
@@ -53,22 +176,34 @@ def save_to_csv(out_path, tracks):
53176 """
54177
55178 with open (out_path , "w" ) as ofile :
56- field_names = ['frame' , 'id' , 'x' , 'y' , 'w' , 'h' , 'score' , 'wx' , 'wy' , 'wz' ]
179+ if fmt == 'motchallenge' :
180+ field_names = ['frame' , 'id' , 'x' , 'y' , 'w' , 'h' , 'score' , 'wx' , 'wy' , 'wz' ]
181+ elif fmt == 'visdrone' :
182+ field_names = ['frame' , 'id' , 'x' , 'y' , 'w' , 'h' , 'score' , 'object_category' , 'truncation' , 'occlusion' ]
183+ else :
184+ raise ValueError ("unknown format type '{}'" .format (fmt ))
57185
58186 odict = csv .DictWriter (ofile , field_names )
59187 id_ = 1
60188 for track in tracks :
61189 for i , bbox in enumerate (track ['bboxes' ]):
62190 row = {'id' : id_ ,
63191 'frame' : track ['start_frame' ] + i ,
64- 'x' : bbox [0 ],
65- 'y' : bbox [1 ],
192+ 'x' : bbox [0 ]+ 1 ,
193+ 'y' : bbox [1 ]+ 1 ,
66194 'w' : bbox [2 ] - bbox [0 ],
67195 'h' : bbox [3 ] - bbox [1 ],
68- 'score' : track ['max_score' ],
69- 'wx' : - 1 ,
70- 'wy' : - 1 ,
71- 'wz' : - 1 }
196+ 'score' : track ['max_score' ]}
197+ if fmt == 'motchallenge' :
198+ row ['wx' ] = - 1
199+ row ['wy' ] = - 1
200+ row ['wz' ] = - 1
201+ elif fmt == 'visdrone' :
202+ row ['object_category' ] = visdrone_classes [track ['class' ]]
203+ row ['truncation' ] = - 1
204+ row ['occlusion' ] = - 1
205+ else :
206+ raise ValueError ("unknown format type '{}'" .format (fmt ))
72207
73208 odict .writerow (row )
74209 id_ += 1
0 commit comments