How do I apply non maximum suppression to multiple images from a dataloader?

1

Hi,

I’m struggling to apply the Non Maximum Suppression (NMS) to multiple images from a dataloader.

I have a function where I apply NMS to one predicted image as seen here, the function defined below:

def apply_nms(orig_prediction, iou_thresh=0.3):
    
    # torchvision returns the indices of the bboxes to keep
    keep = torchvision.ops.nms(orig_prediction['boxes'], orig_prediction['scores'], iou_thresh)
    
    final_prediction = orig_prediction
    final_prediction['boxes'] = final_prediction['boxes'][keep]
    final_prediction['scores'] = final_prediction['scores'][keep]
    final_prediction['labels'] = final_prediction['labels'][keep]
    
    return final_prediction

Here I take one image from the validation set and infer some predictions, where I then apply NMS to that one image:

# one image from the test set
img, target = valid_dataset[2]
# put the model in evaluation mode
model.eval()
with torch.no_grad():
    prediction = model([img.to(device)])[0]

nms_prediction = apply_nms(prediction, iou_thresh=0.1)

How can I do this for multiple images (i.e. a batch of images)?

I currently have this set up:

model.eval()
with torch.no_grad():
  for images, targets in valid_data_loader:
    images = list(img.to(device) for img in images)
    predictions = model(images)

However, am struggling to understand how I can apply NMS to multiple images without using for loops…