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 # put the model in evaluation mode model.eval() with torch.no_grad(): prediction = model([img.to(device)]) 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…