glenn-jocher commited on
Commit
5a79b5f
·
1 Parent(s): 4f44aaf
Files changed (1) hide show
  1. utils/utils.py +3 -1
utils/utils.py CHANGED
@@ -473,6 +473,7 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=T
473
  min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
474
  max_det = 300 # maximum number of detections per image
475
  time_limit = 10.0 # seconds to quit after
 
476
 
477
  t = time.time()
478
  nc = prediction[0].shape[1] - 5 # number of classes
@@ -528,7 +529,8 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, multi_label=T
528
  iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
529
  weights = iou * scores[None] # box weights
530
  x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
531
- # i = i[iou.sum(1) > 1] # require redundancy
 
532
  except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139
533
  print(x, i, x.shape, i.shape)
534
  pass
 
473
  min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
474
  max_det = 300 # maximum number of detections per image
475
  time_limit = 10.0 # seconds to quit after
476
+ redundant = conf_thres == 0.001 # require redundant detections
477
 
478
  t = time.time()
479
  nc = prediction[0].shape[1] - 5 # number of classes
 
529
  iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
530
  weights = iou * scores[None] # box weights
531
  x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
532
+ if redundant:
533
+ i = i[iou.sum(1) > 1] # require redundancy
534
  except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139
535
  print(x, i, x.shape, i.shape)
536
  pass