This commit is contained in:
Glenn Jocher 2018-11-23 15:34:49 +01:00
parent 6e825acb72
commit bf66656b4e
1 changed files with 1 additions and 1 deletions

View File

@ -306,7 +306,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
# thresh = 0.85 # thresh = 0.85
thresh = nms_thres thresh = nms_thres
a = pred.clone() a = pred.clone()
_, indices = torch.sort(-a[:, 4], 0) # sort best to worst _, indices = torch.sort(-a[:, 4], 0) # sort best to worst
a = a[indices] a = a[indices]
radius = 30 # area to search for cross-class ious radius = 30 # area to search for cross-class ious
for i in range(len(a)): for i in range(len(a)):