Browse Source

Adjust NMS time limit warning to batch size (#7156)

modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
26bfd44465
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 4 deletions
  1. +7
    -4
      utils/general.py

+ 7
- 4
utils/general.py View File

list of detections, on (n,6) tensor per image [xyxy, conf, cls] list of detections, on (n,6) tensor per image [xyxy, conf, cls]
""" """


bs = prediction.shape[0] # batch size
nc = prediction.shape[2] - 5 # number of classes nc = prediction.shape[2] - 5 # number of classes
xc = prediction[..., 4] > conf_thres # candidates xc = prediction[..., 4] > conf_thres # candidates


# Settings # Settings
min_wh, max_wh = 2, 7680 # (pixels) minimum and maximum box width and height min_wh, max_wh = 2, 7680 # (pixels) minimum and maximum box width and height
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms() max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
time_limit = 10.0 # seconds to quit after
time_limit = 0.030 * bs # seconds to quit after
redundant = True # require redundant detections redundant = True # require redundant detections
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img) multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
merge = False # use merge-NMS merge = False # use merge-NMS


t = time.time()
output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
t, warn_time = time.time(), True
output = [torch.zeros((0, 6), device=prediction.device)] * bs
for xi, x in enumerate(prediction): # image index, image inference for xi, x in enumerate(prediction): # image index, image inference
# Apply constraints # 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


output[xi] = x[i] output[xi] = x[i]
if (time.time() - t) > time_limit: if (time.time() - t) > time_limit:
LOGGER.warning(f'WARNING: NMS time limit {time_limit}s exceeded')
if warn_time:
LOGGER.warning(f'WARNING: NMS time limit {time_limit:3f}s exceeded')
warn_time = False
break # time limit exceeded break # time limit exceeded


return output return output

Loading…
Cancel
Save