Plot `torch.Tensor` using OpenCV

I have a pytorch tensor, let’s say images, of type <class 'torch.Tensor'> and of size torch.Size([32, 3, 300, 300]), so that images[i, :, :, :] represents the i-th out of 32 rgb 300x300 images.

I would like to plot some of these images, let’s say the first one, i.e., images[0, :, :, :] using opencv (cv2.imshow()). Could you help? Thank you!

1 Like

You can call images[0].numpy() to get the numpy array, which can be visualized using OpenCV.

1 Like

Thanks for your response. I tried that but I got an opencv error (Source image must have 1, 3 or 4 channels in function cvConvertImage)… It seemed to me that it would be trivial, but it isn’t…

Sorry, my bad. You have to permute the axis:

images[0].numpy().transpose(1, 2, 0)
# or
images[0].permute(1, 2, 0).numpy()
8 Likes

I see. Thank you very much, it worked!

Why can’t we use torch tensors directly with Open cv without converting them to numpy? I run cv2.kmeans on large images and they’re very slow. Wondering if being able to run them on Tensors would be faster.

1 Like

after converting your torch tensor back to opencv ndarray, if you do an imshow the image will appear slightly darker due to standard normalization.

def inverse_normalize(tensor, mean, std):
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)
    return tensor

input = inverse_normalize(tensor=input, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))

you should inverse normalize your torch tensor before converting to numpy array if the image colours matter to you

1 Like

I think opencv-python package support CPU-only. So, we need to change cuda tensor to cpu.