How to Outputting tensor gpu format as imread?

vision
1

mask_pred = net(imgs)

It is mapped to values of 0.96, 0.006, and so on. (because of normalization: x/ 255)

I’d like to see if this format was put in imshow() to see if it was well segmented

so I try…

imgs # size(1,1,640,640) = size(batch_size, channel, width, height)

    mask_pred = net(imgs)[0]

    mask_pred = (mask_pred >0.5).float()

    np_img = mask_pred.cpu().detach().numpy() # GPU Tensor -> CPU Tensor -> numpy

    np_img = np.transpose(np_img, (1, 2, 0)) # <class 'tuple'>: (640,640,1), dtype=float32

    plt.imshow(np_img)

TypeError: Invalid dimensions for image data

how to solve it?

2

Since your image has only 1 channel, then you should remove the last channel. So instead of having shape (640, 640, 1) you should have (640, 640):

np_img = np_img[:, :, 0]
plt.imshow(np_img)