Issue applying mask transform to 4d tensor

Is there an existing issue for this?

• I have searched the existing issues

Problem summary

Hello, I am trying to apply a brain mask (1,l,w,d) to a 4d series (DSC MRI sequence) (120,l,w,d), however, when I run the code I get an error message saying that the transform cannot be applied because there is a mismatch in the channel dimension between the label and the image. I guess this issue is somewhat related to a previous issue "Support 4D tensors as torchio.Image #213 ", but I was wondering if there was a straightforward way to implement this transformation?

Code for reproduction

subj_dir=data_dir+'/QIN-BRAIN-DSC-01-46'   
subj_dict = {'patientID':'QIN-BRAIN-DSC-01-46', 'brain_mask': tio.LabelMap(subj_dir+'/Brain_mask.nii'),'DSC_sequence':tio.ScalarImage(subj_dir+'/dsc.nii'), 'perf_map':tio.ScalarImage(subj_dir+'/rBF_LC.nii')}
    
subject = tio.Subject(subj_dict)
subject

subject.DSC_sequence.shape

mask = tio.Mask(masking_method='brain_mask')  # Use "brain" image to mask DSC_sequence and perf_map images. DSC_sequence is a (120, 256, 256, 13) tensor
transformed = mask(subject)
transformed

Actual outcome

I get an error message when I try to train my model.

Error messages

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Input In [6], in <cell line: 2>()
      1 mask = tio.Mask(masking_method='brain_mask')  # Use "brain" image to mask
----> 2 transformed = mask(subject)
      3 transformed

File ~/miniconda3/envs/pytorchenv/lib/python3.9/site-packages/torchio/transforms/transform.py:140, in Transform.__call__(self, data)
    138     subject = copy.copy(subject)
    139 with np.errstate(all='raise', under='ignore'):
--> 140     transformed = self.apply_transform(subject)
    141 if self.keep is not None:
    142     for name, image in images_to_keep.items():

File ~/miniconda3/envs/pytorchenv/lib/python3.9/site-packages/torchio/transforms/preprocessing/intensity/mask.py:62, in Mask.apply_transform(self, subject)
     55 for image in self.get_images(subject):
     56     mask_data = self.get_mask_from_masking_method(
     57         self.masking_method,
     58         subject,
     59         image.data,
     60         self.masking_labels,
     61     )
---> 62     self.apply_masking(image, mask_data)
     63 return subject

File ~/miniconda3/envs/pytorchenv/lib/python3.9/site-packages/torchio/transforms/preprocessing/intensity/mask.py:66, in Mask.apply_masking(self, image, mask_data)
     65 def apply_masking(self, image: LabelMap, mask_data: torch.Tensor) -> None:
---> 66     masked = mask(image.data, mask_data, self.outside_value)
     67     image.set_data(masked)

File ~/miniconda3/envs/pytorchenv/lib/python3.9/site-packages/torchio/transforms/preprocessing/intensity/mask.py:77, in mask(tensor, mask, outside_value)
     75 array = tensor.clone().numpy()
     76 mask = mask.numpy()
---> 77 array[~mask] = outside_value
     78 return torch.as_tensor(array)

IndexError: boolean index did not match indexed array along dimension 0; dimension is 120 but corresponding boolean dimension is 1

Expected outcome

I should be able to obtain a DSC_sequence image brain-masked in all 120 channels.

System info

Platform:   Linux-5.3.0-59-generic-x86_64-with-glibc2.27
TorchIO:    0.18.75
PyTorch:    1.11.0
SimpleITK:  2.1.1 (ITK 5.2)
NumPy:      1.21.2
Python:     3.9.7 (default, Sep 16 2021, 13:09:58) 
[GCC 7.5.0]

[fepegar]

Resolved in v0.18.79. Thank you both for reporting and for your insights.

[cbri92]

Thank you @fepegar