I don’t think there is an problem with reshaping image.
I printed out the shape just before the imshow() and It was 3 dimension.
r = random.randint(0, 3608)
gen = BasicDataset('/home/intern/Desktop/YH/Brain_MRI/BrainMRI_train/MRI/MRI/', '/home/intern/Desktop/YH/Brain_MRI/BrainMRI_train/mask/mask/')
x, y = gen.__getitem__(r)
img = x
mask = y["masks"]

img = np.transpose(img, (1, 2, 0)) #makes 256, 256, 3
mask = mask.numpy()
mask = np.transpose(mask, (1, 2, 0)) #makes 256, 256, 1

print("img shape :", img.shape, "\nmask shape :", mask.shape, "\n")
fig = plt.figure()
fig.subplots_adjust(hspace=0.4, wspace=0.4)
ax = fig.add_subplot(1, 2, 1)
ax.imshow(img)
ax = fig.add_subplot(1, 2, 2)
ax.imshow(np.reshape(mask,(256, 256)), cmap = "gray")

out
img shape : (256, 256, 3) 
mask shape : (256, 256, 1)


imagesds011044×478 113 KB

mine is working good