tuansunday05 commited on
Commit
b28d005
·
verified ·
1 Parent(s): b917f56

Update detect_strongsort.py

Browse files
Files changed (1) hide show
  1. detect_strongsort.py +1 -0
detect_strongsort.py CHANGED
@@ -185,6 +185,7 @@ def run_strongsort(
185
 
186
  # Apply NMS
187
  with dt[2]:
 
188
  pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
189
  sdt[2] += time_sync() - t3
190
 
 
185
 
186
  # Apply NMS
187
  with dt[2]:
188
+ pred = pred[0][1] if isinstance(pred[0], list) else pred[0] # single model or ensemble
189
  pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
190
  sdt[2] += time_sync() - t3
191