This commit is contained in:
Glenn Jocher 2019-04-24 17:13:36 +02:00
parent aa2df1eda7
commit fbf0014cd6
1 changed files with 5 additions and 5 deletions

View File

@ -376,12 +376,12 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):
nms_style = 'MERGE' # 'OR' (default), 'AND', 'MERGE' (experimental) nms_style = 'MERGE' # 'OR' (default), 'AND', 'MERGE' (experimental)
for c in pred[:, -1].unique(): for c in pred[:, -1].unique():
dc = pred[pred[:, -1] == c] # select class c dc = pred[pred[:, -1] == c] # select class c
dc = dc[:min(len(dc), 100)] # limit to first 100 boxes: https://github.com/ultralytics/yolov3/issues/117 n = len(dc)
if n == 1:
# No NMS required if only 1 prediction det_max.append(dc) # No NMS required if only 1 prediction
if len(dc) == 1:
det_max.append(dc)
continue continue
elif n > 100:
dc = dc[:100] # limit to first 100 boxes: https://github.com/ultralytics/yolov3/issues/117
# Non-maximum suppression # Non-maximum suppression
if nms_style == 'OR': # default if nms_style == 'OR': # default