|
|
@@ -703,7 +703,7 @@ def clip_coords(boxes, shape): |
|
|
|
|
|
|
|
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, |
|
|
|
labels=(), max_det=300): |
|
|
|
"""Runs Non-Maximum Suppression (NMS) on inference results |
|
|
|
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping bounding boxes |
|
|
|
|
|
|
|
Returns: |
|
|
|
list of detections, on (n,6) tensor per image [xyxy, conf, cls] |
|
|
@@ -718,18 +718,19 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non |
|
|
|
assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0' |
|
|
|
|
|
|
|
# Settings |
|
|
|
min_wh, max_wh = 2, 7680 # (pixels) minimum and maximum box width and height |
|
|
|
# min_wh = 2 # (pixels) minimum box width and height |
|
|
|
max_wh = 7680 # (pixels) maximum box width and height |
|
|
|
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms() |
|
|
|
time_limit = 0.030 * bs # seconds to quit after |
|
|
|
redundant = True # require redundant detections |
|
|
|
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img) |
|
|
|
merge = False # use merge-NMS |
|
|
|
|
|
|
|
t, warn_time = time.time(), True |
|
|
|
t = time.time() |
|
|
|
output = [torch.zeros((0, 6), device=prediction.device)] * bs |
|
|
|
for xi, x in enumerate(prediction): # image index, image inference |
|
|
|
# Apply constraints |
|
|
|
x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height |
|
|
|
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height |
|
|
|
x = x[xc[xi]] # confidence |
|
|
|
|
|
|
|
# Cat apriori labels if autolabelling |
|
|
@@ -790,9 +791,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non |
|
|
|
|
|
|
|
output[xi] = x[i] |
|
|
|
if (time.time() - t) > time_limit: |
|
|
|
if warn_time: |
|
|
|
LOGGER.warning(f'WARNING: NMS time limit {time_limit:3f}s exceeded') |
|
|
|
warn_time = False |
|
|
|
LOGGER.warning(f'WARNING: NMS time limit {time_limit:.3f}s exceeded') |
|
|
|
break # time limit exceeded |
|
|
|
|
|
|
|
return output |