updates
This commit is contained in:
parent
74b57500c7
commit
a0067ac8fb
|
@ -442,7 +442,7 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):
|
||||||
(x1, y1, x2, y2, object_conf, class_conf, class)
|
(x1, y1, x2, y2, object_conf, class_conf, class)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
min_wh = 2 # (pixels) minimum box width and height
|
min_wh, max_wh = 2, 30000 # (pixels) minimum and maximium box width and height
|
||||||
|
|
||||||
output = [None] * len(prediction)
|
output = [None] * len(prediction)
|
||||||
for image_i, pred in enumerate(prediction):
|
for image_i, pred in enumerate(prediction):
|
||||||
|
@ -470,7 +470,8 @@ def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.5):
|
||||||
# pred[class_pred != 2, 4] = 0.0
|
# pred[class_pred != 2, 4] = 0.0
|
||||||
|
|
||||||
# Select only suitable predictions
|
# Select only suitable predictions
|
||||||
i = (pred[:, 4] > conf_thres) & (pred[:, 2:4] > min_wh).all(1) & torch.isfinite(pred).all(1)
|
i = (pred[:, 4] > conf_thres) & (pred[:, 2:4] > min_wh).all(1) & (pred[:, 2:4] < max_wh).all(1) & \
|
||||||
|
torch.isfinite(pred).all(1)
|
||||||
pred = pred[i]
|
pred = pred[i]
|
||||||
|
|
||||||
# If none are remaining => process next image
|
# If none are remaining => process next image
|
||||||
|
|
Loading…
Reference in New Issue