LabelMap not tranformed in tio.Subject when label shape is 1xHxWx1

[chr-wei]

🐛 Describe the bug

Situation: In a dataloader a tio composition of intensity and spatial augmentation is performed. When using 3D samples the tio.ScalarImage and tio.LabelMap are transformed fine, when using 2D samples the tio.LabelMap is not spatially transformed. This only happens when within the tio.Compose a intensity transform is applied along with spatial transform (either prior or after spatial transformation).

import torchio as tio

if self.yield_2d:            
    self.spatial_transform = tio.OneOf({
        tio.transforms.RandomAffine(.05): 0.8,
        tio.transforms.RandomElasticDeformation(num_control_points=7, max_displacement=(7.5,7.5,0)): 0.2,
        },
        p=0.75,
    )
else:
    self.spatial_transform = tio.OneOf({
        tio.transforms.RandomAffine(.05): 0.8,
        tio.transforms.RandomElasticDeformation(num_control_points=7, max_displacement=7.5): 0.2,
        },
        p=0.75,
    )

self.transform = tio.Compose([
    self.intensity_transform,
    self.spatial_transform
])

# 2D image are dim=1xHxWx1 after unsqueezing
# 3D image are dim=1xDxHxW after unsqueezing

if yield_2d:
    image = image.unsqueeze(0).unsqueeze(-1)
    label = label.unsqueeze(0).unsqueeze(-1)
else:
    image = image.unsqueeze(0)
    label = label.unsqueeze(0)

subject = tio.Subject(
    image=tio.ScalarImage(tensor=image),  
    label=tio.LabelMap(tensor=label)
)
subject = self.spatial_transform(subject)

3D images are transformed fine:

2D images are not:

Versions

Collecting environment information... PyTorch version: 1.9.1+cu102 Is debug build: False CUDA used to build PyTorch: 10.2 ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.7 LTS (x86_64) GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609 Clang version: 3.8.0-2ubuntu4 (tags/RELEASE_380/final) CMake version: version 3.5.1 Libc version: glibc-2.23

Python version: 3.9.5 (default, Jun 4 2021, 12:28:51) [GCC 7.5.0] (64-bit runtime) Python platform: Linux-4.4.0-210-generic-x86_64-with-glibc2.23 Is CUDA available: True CUDA runtime version: 7.5.17 GPU models and configuration: GPU 0: Tesla T4 GPU 1: Tesla T4 GPU 2: GeForce RTX 2080 Ti GPU 3: Quadro RTX 8000 GPU 4: GeForce RTX 2080 Ti GPU 5: GeForce RTX 2080 Ti GPU 6: GeForce RTX 2080 Ti

Nvidia driver version: 455.45.01 cuDNN version: Probably one of the following: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.0 /usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.0 /usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.0 /usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.0 /usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.0 /usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.0 /usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.0 /usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.0 HIP runtime version: N/A MIOpen runtime version: N/A

Versions of relevant libraries: [pip3] numpy==1.21.1 [pip3] torch==1.9.1 [pip3] torchio==0.18.63 [pip3] torchvision==0.10.1 [conda] Could not collect

[datumbox]

@rap1ide Should this be raised to the torchio repo directly? It seems to me that's probably the best place to kick off the investigation.

[fepegar]

@chr-wei, please ask in https://github.com/fepegar/torchio as suggested by @datumbox.