No bouding box displayed with draw_bouding_boxes from pytorch

21 Views Asked by At

I am performing inference on an img for Object detection, I try to run the below code, it correctly runs and displays the image but without any bounding boxes, The bounding box coordinates are correctly printing with the scores and labels as well.

mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225]
def eval_img(model_path, img_no):
    model = torch.load(model_path)
    model.eval()
    
    feature, label = next(iter(val_loader))
    with torch.no_grad():
        img = [feature[img_no].to(device)]
        output = model(img)
        pred = output[0]
    
    print(pred)
    img = feature[img_no].permute(1,2,0).numpy()
    img = (mean+img*std)
    new_img = torch.tensor(img*255, dtype=torch.uint8).permute(2,0,1)
    
    iou_thresh = 0.10
    keep = torchvision.ops.nms(pred['boxes'], pred['scores'], iou_thresh)
    pred['boxes'] = pred['boxes'][keep]
    pred['scores'] = pred['scores'][keep]
    pred['labels'] = pred['labels'][keep]

    plt.imshow(torchvision.utils.draw_bounding_boxes(new_img, pred['boxes'][pred['scores']>0.9], width=4)).permute(1,2,0)

Following is the output from printing the pred variable -

{'boxes': tensor([[415.9670, 225.6164, 416.0000, 248.0225],
        [415.8290, 223.3006, 415.9943, 243.7530]], device='cuda:0'), 'labels': tensor([1, 1], device='cuda:0'), 'scores': tensor([0.3576, 0.0766], device='cuda:0')}

The boxes and the scores values are correctly being printed, but no bounding box is displayed

I initially varied the IoU threshold values for NMS and tried this on couple of pictures, but to no avail

0

There are 0 best solutions below