Loading and correctly displaying an image dataset using pytorch Dataloader
I'm trying to load a custom dataset for training a neural network, but before I load them in, I would like to verify that they've been loaded correctly. So far it looks like they are not being loaded correctly, but I can't figure out what gives the images the format that they get.

This is the code that I'm loading the images with, and then displaying them.

f, axarr = plt.subplots(2,2, figsize=(20,20))

def load_dataset():
    data_path = 'processedData/HE/train/'
    train_dataset = torchvision.datasets.ImageFolder(
        root=data_path,
        transform=torchvision.transforms.ToTensor()
    )
    train_loader = DataLoader(
        train_dataset, batch_size=64
    )
    return train_loader
x_train = load_dataset()

datathing = next(iter(x_train))

for i, ax in enumerate(axarr.flat):
    ax.imshow(datathing[0][i].view(128,128,3))
    ax.axis('off')
plt.show()


When running this with the images, the output looks like this.

It is suppose to look something like these images

I have been trying with different image datasets, but all the sets return the same format, so my question is:


How to load images and then displaying them in their real format using pytorch's dataloader?


The .view(128, 128, 3) is messing up with the images.
As you can read in the documentation of the transformation .ToTensor(...):

[...]
Converts a PIL Image or numpy.ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]
[...]

That is, the channel dimension is moved from the last to the first dimension. You can see that in the source code:
# ...
img = torch.from_numpy(pic.transpose((2, 0, 1)))
# ...

Therefore, you cannot simply call .view(...); you have to transpose it back. In PyTorch, you can use the .permute(...) function for that. Something like this:
ax.imshow(datathing[0][i].permute(1, 2, 0))

