ToTensorV2 assumes image dimension as 3

🐛 Bug

albumentations.pytorch.ToTensorV2 assumes that input image has 3 dimensions since it calls

torch.from_numpy(img.transpose(2, 0, 1))

The transpose operation throws a ValueError axes don’t match array when any image other than 3 dimensional is provided.

To Reproduce

Steps to reproduce the behavior:

1. Create a PyTorch loader, similar like here: https://github.com/albu/albumentations/blob/master/notebooks/migrating_from_torchvision_to_albumentations.ipynb
2. Apply the transformation in getitem self._transforms(image=image)[‘image’]

Expected behavior

Transpose and ToTensor calls should be separated. It blocks users from using grayscale or any high dimensional input data.

Environment

• Albumentations version (e.g., 0.1.8): 0.4.3
• Python version (e.g., 3.7): 3.6.8
• OS (e.g., Linux): CentOS 7
• How you installed albumentations (conda, pip, source): pip

[AntixK]

Yes, I get that but a function named ToTensor must do… just that… convert to tensor. Adding extra (undocumented) functionality only leads to confusion.

I already have a pre-processed dataset in CHW format and I would like to use that for TTA. Using the current ToTensorV2 unnecessarily requires me to transpose the channels just to pass to the function. Of course, I can have a workaround by having my own ToTensor, but doesn’t that forego the idea of ToTensorV2 in the first place?

[Dipet]

If you need only to convert image to tensor you has function torch.from_numpy. This transform is necessary to convert images from the library format to the torch image format. And torch work with images in CHW, but library works in HWC, so we need transpose.