How can I apply non-maximum suppression (NMS) to a batch of images?

5

Thanks for the suggestions!

I did try it with a batch size of 1 and it does work which is useful, I’m not sure how efficient this is but it’s better than not applying NMS!

# batch size of 1 in dataloader

cpu_device = torch.device("cpu")
model.eval()
with torch.no_grad():
  for images, targets in valid_data_loader:
    images = list(img.to(device) for img in images)
    outputs = model(images)
    outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
    predictions = apply_nms(outputs[0], iou_thresh=0.3)

My end goal is to then find the f1-score. I don’t know if this warrants another question on the forum, but I’m not really sure how I can extract true positives, false positives etc. from the outputs…