@@ -193,29 +193,34 @@ def __call__(self, scores, raw_boxes):
193
193
top_k = self .keep_top_k , )
194
194
picked_box_probs .append (box_probs )
195
195
picked_labels .extend ([class_index ] * box_probs .shape [0 ])
196
- if not picked_box_probs :
197
- return np .array ([]), np .array ([]), np .array ([])
198
- picked_box_probs = np .concatenate (picked_box_probs )
199
-
200
- # resize output boxes
201
- picked_box_probs [:, :4 ] = self .warp_boxes (picked_box_probs [:, :4 ],
202
- self .ori_shape [batch_id ])
203
- im_scale = np .concatenate ([
204
- self .scale_factor [batch_id ][::- 1 ],
205
- self .scale_factor [batch_id ][::- 1 ]
206
- ])
207
- picked_box_probs [:, :4 ] /= im_scale
208
- # clas score box
209
- out_boxes_list .append (
210
- np .concatenate (
211
- [
212
- np .expand_dims (
213
- np .array (picked_labels ), axis = - 1 ), np .expand_dims (
214
- picked_box_probs [:, 4 ], axis = - 1 ),
215
- picked_box_probs [:, :4 ]
216
- ],
217
- axis = 1 ))
218
- out_boxes_num .append (len (picked_labels ))
196
+
197
+ if len (picked_box_probs ) == 0 :
198
+ out_boxes_list .append (np .empty ((0 , 4 )))
199
+ out_boxes_num .append (0 )
200
+
201
+ else :
202
+ picked_box_probs = np .concatenate (picked_box_probs )
203
+
204
+ # resize output boxes
205
+ picked_box_probs [:, :4 ] = self .warp_boxes (
206
+ picked_box_probs [:, :4 ], self .ori_shape [batch_id ])
207
+ im_scale = np .concatenate ([
208
+ self .scale_factor [batch_id ][::- 1 ],
209
+ self .scale_factor [batch_id ][::- 1 ]
210
+ ])
211
+ picked_box_probs [:, :4 ] /= im_scale
212
+ # clas score box
213
+ out_boxes_list .append (
214
+ np .concatenate (
215
+ [
216
+ np .expand_dims (
217
+ np .array (picked_labels ),
218
+ axis = - 1 ), np .expand_dims (
219
+ picked_box_probs [:, 4 ], axis = - 1 ),
220
+ picked_box_probs [:, :4 ]
221
+ ],
222
+ axis = 1 ))
223
+ out_boxes_num .append (len (picked_labels ))
219
224
220
225
out_boxes_list = np .concatenate (out_boxes_list , axis = 0 )
221
226
out_boxes_num = np .asarray (out_boxes_num ).astype (np .int32 )
0 commit comments