Inconsistent shapes for diffusion MRI

Hi all, I am willing to use Torchio for diffusion MRI and this data is 4D. My plan is to use TorchIO for a patch-based regression where both the input and the output are diffusion MR images. To sample the patches, I am planning to use the intra-cranial mask as a sampling prior.

I ran the following code to try to apply a rotation transform to a subject and I have a ValueError: Images in sample have inconsistent shapes. I would appreciate any help with that issue.

import torchio
from torchio.transforms import RandomAffine


image = '/xxx/aylj.nii.gz'
label = '/xxxx/aylj.nii.gz'
sample = '/xxxx/maskaylj.nii.gz'
subject = torchio.Subject(dMRI=torchio.Image(image, type=torchio.INTENSITY),
                          label=torchio.Image(image, type=torchio.INTENSITY),
                          sampling_map=torchio.Image(sample, type=torchio.SAMPLING_MAP))


transform = RandomAffine(degrees=(-10, 10),
                         isotropic=True)

transformed = transform(subject)

patch_size = 64
sampler = torchio.data.WeightedSampler(patch_size, probability_map='sampling_map')
sample = torchio.ImagesDataset([transformed])[0]

Traceback (most recent call last):
  File "test_torchio.py", line 17, in <module>
    transformed = transform(subject)
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/transforms/augmentation/random_transform.py", line 33, in __call__
    return super().__call__(sample)
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/transforms/transform.py", line 62, in __call__
    transformed = self.apply_transform(sample)
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/transforms/augmentation/spatial/random_affine.py", line 154, in apply_transform
    sample.check_consistent_shape()
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/data/subject.py", line 122, in check_consistent_shape
    raise ValueError(message)
ValueError: Images in sample have inconsistent shapes:
{'dMRI': (1, 128, 128, 69, 15),
 'label': (1, 128, 128, 69, 15),
 'sampling_map': (1, 128, 128, 69)}

Additionally, If I remove the sample from the Subject, the following error is raised:

Traceback (most recent call last):
  File "test_torchio.py", line 15, in <module>
    transformed = transform(subject)
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/transforms/augmentation/random_transform.py", line 33, in __call__
    return super().__call__(sample)
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/transforms/transform.py", line 62, in __call__
    transformed = self.apply_transform(sample)
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/transforms/augmentation/spatial/random_affine.py", line 184, in apply_transform
    center_lps=center,
  File "/home/ol18/miniconda3/envs/tract_uncertainty/lib/python3.6/site-packages/torchio/transforms/augmentation/spatial/random_affine.py", line 204, in apply_affine_transform
    assert tensor.ndim == 4
AssertionError

[OeslleLucena]

@fepegar sure! I would do that 😃

[OeslleLucena]

I guess that might work temporarily. I will see if I can write my scripts based on that workaround. Thanks @fepegar