Spacingd changes the output shape unexpectedly

[josiglo]

Describe the bug
Assume I have three images A, B and C, which I want to transform using one Spacingd instance.

Images A and B already have the desired target pixdim, I thus expect the application of Spacingd to not have an effect. Image C does not have the desired target pixdim, I thus expect the application of Spacingd to have an effect.

In case I pass the image keys in order [A, B, C] or [C, A, B] to Spacingd, the transformed images A and B do not differ, which is what I expect.

In case I pass the image keys in order [A, C, B] to Spacingd, the transformed images A and B differ although I expect them not to. The output shape of image B has changed although I expect it not to.

Please see the attached code for details.

To Reproduce
Make sure to install Monai v1.2.0 and run the following code.

"""
Potential bug in Spacingd (v1.2.0).

Summary: The order of keys passed to Spacingd seems to make a difference.

We will pass three keys NII_1MM_A, NII_1MM_B and NII_2MM as well as a target spacing of 1mm^3 to
Spacingd. NII_1MM_A and NII_1MM_B are identical NIfTIs that are already in 1mm spacing. We
thus expect these two NIfTIs to be identical and unchanged after applying Spacingd.

We run three tests that differ in the order of keys passed to Spacingd and expect results for all
three tests to be identicial, which is not the case.
"""
from tempfile import NamedTemporaryFile

import nibabel as nib
import numpy as np
import torch
from monai.config.type_definitions import KeysCollection
from monai.transforms.compose import Compose
from monai.transforms.io.dictionary import LoadImaged
from monai.transforms.spatial.dictionary import Spacingd

NII_1MM_A = "nii_1mm_a"
NII_1MM_B = "nii_1mm_b"
NII_2MM = "nii_2mm"

TARGET_SPACING = (1, 1, 1)

AFFINE_1MM_SPACING = np.eye(4)
AFFINE_2MM_SPACING = np.diag([2, 2, 2, 1])

nii_1mm = nib.Nifti1Image(np.ones((5, 5, 5)), affine=AFFINE_1MM_SPACING)
nii_2mm = nib.Nifti1Image(np.ones((5, 5, 5)), affine=AFFINE_2MM_SPACING)


def run_test(keys: KeysCollection):
    transforms = Compose([
        LoadImaged(keys),
        Spacingd(keys, pixdim=TARGET_SPACING),
    ])

    with NamedTemporaryFile(suffix=".nii") as f1, NamedTemporaryFile(
            suffix=".nii") as f2, NamedTemporaryFile(suffix=".nii") as f3:
        f1.write(nii_1mm.to_bytes())
        f2.write(nii_2mm.to_bytes())
        f3.write(nii_1mm.to_bytes())
        f1.flush()
        f2.flush()
        f3.flush()
        data = {NII_1MM_A: f1.name, NII_2MM: f2.name, NII_1MM_B: f3.name}
        transformed = transforms(data)

    print("Tensors NII_1MM_A and NII_1MM_B are expected to be equal. Are they?",
          torch.equal(transformed[NII_1MM_A], transformed[NII_1MM_B]))
    print("Expected shape of tensor NII_1MM_A is (5, 5, 5). Actual shape:",
          transformed[NII_1MM_A].shape)
    print("Expected shape of tensor NII_1MM_B is (5, 5, 5). Actual shape:",
          transformed[NII_1MM_B].shape)
    print("--")


# NII_2MM before other keys -> Working as expected
run_test([NII_2MM, NII_1MM_A, NII_1MM_B])

# NII_2MM after other keys -> Working as expected
run_test([NII_1MM_A, NII_1MM_B, NII_2MM])

# NII_2MM in between  other keys -> Not working as expected
run_test([NII_1MM_A, NII_2MM, NII_1MM_B])

Expected behavior
I expect identical results from Spacingd, regardless of the order of keys passed to it.

Environment

================================
Printing MONAI config...
================================
MONAI version: 1.2.0+92.g40048d73
Numpy version: 1.23.5
Pytorch version: 2.0.1+cu117
MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False
MONAI rev id: 40048d732f0d9bfb975edda0cb24e51464d18ae0
MONAI __file__: /home/airamed.local/jsi/.pyenv/versions/datascience/lib/python3.8/site-packages/monai/__init__.py

Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
ITK version: 5.4.0
Nibabel version: 5.1.0
scikit-image version: 0.21.0
scipy version: 1.10.1
Pillow version: 10.0.0
Tensorboard version: 2.12.3
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: NOT INSTALLED or UNKNOWN VERSION.
tqdm version: 4.66.1
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: 5.9.5
pandas version: 2.0.3
einops version: NOT INSTALLED or UNKNOWN VERSION.
transformers version: NOT INSTALLED or UNKNOWN VERSION.
mlflow version: NOT INSTALLED or UNKNOWN VERSION.
pynrrd version: NOT INSTALLED or UNKNOWN VERSION.
clearml version: NOT INSTALLED or UNKNOWN VERSION.

For details about installing the optional dependencies, please visit:
    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies


================================
Printing system config...
================================
System: Linux
Linux version: Ubuntu 22.04.3 LTS
Platform: Linux-6.2.0-26-generic-x86_64-with-glibc2.35
Processor: x86_64
Machine: x86_64
Python version: 3.8.12
Process name: python
Command: ['/home/airamed.local/jsi/.pyenv/versions/datascience/bin/python', '-c', 'import monai; monai.config.print_debug_info()']
Open files: [popenfile(path='/home/airamed.local/jsi/.vscode-server/data/logs/20230829T095406/ptyhost.log', fd=19, position=4270, mode='a', flags=33793), popenfile(path='/home/airamed.local/jsi/.vscode-server/data/logs/20230829T095406/remoteagent.log', fd=20, position=12299, mode='a', flags=33793), popenfile(path='/home/airamed.local/jsi/.vscode-server/bin/6261075646f055b99068d3688932416f2346dd3b/vscode-remote-lock.jsi.6261075646f055b99068d3688932416f2346dd3b', fd=99, position=0, mode='w', flags=32769)]
Num physical CPUs: 6
Num logical CPUs: 12
Num usable CPUs: 12
CPU usage (%): [4.2, 3.3, 3.3, 2.9, 3.6, 85.1, 3.9, 3.2, 2.9, 3.6, 3.3, 17.6]
CPU freq. (MHz): 1611
Load avg. in last 1, 5, 15 mins (%): [2.5, 2.7, 2.9]
Disk usage (%): 94.3
Avg. sensor temp. (Celsius): UNKNOWN for given OS
Total physical memory (GB): 31.2
Available memory (GB): 17.9
Used memory (GB): 12.9

================================
Printing GPU config...
================================
Num GPUs: 1
Has CUDA: True
CUDA version: 11.7
cuDNN enabled: True
NVIDIA_TF32_OVERRIDE: None
TORCH_ALLOW_TF32_CUBLAS_OVERRIDE: None
cuDNN version: 8500
Current device: 0
Library compiled for CUDA architectures: ['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75', 'sm_80', 'sm_86']
GPU 0 Name: Quadro K2200
GPU 0 Is integrated: False
GPU 0 Is multi GPU board: False
GPU 0 Multi processor count: 5
GPU 0 Total memory (GB): 3.9
GPU 0 CUDA capability (maj.min): 5.0

Additional context
I also tried the current dev version of Monai, but experienced the same thing. The output shape is not changed by Spacingd for Monai versions <=1.1.0.

Could be related to #5935 but from what I saw the PR that resulted from it made it into v1.2.0.

[KumoLiu]

Hi @josiglo, looks like a bug introduced by #6126.
It will always use the affine from the first key.

MONAI/monai/transforms/spatial/dictionary.py

Line 505 in 40048d7

_init_shape, _pixdim = d[key].peek_pending_shape(), d[key].pixdim

cc @wyli

[wyli]

thanks for looking into this issue @josiglo @KumoLiu, I think the root cause is from

MONAI/monai/transforms/spatial/dictionary.py

Line 520 in 40048d7

output_shape_k = d[key].peek_pending_shape() if isinstance(d[key], MetaTensor) else d[key].shape[1:]

output_shape_k is always updated to use the most recent key, which is not consistent with the way _init_shape and _pixdim are updated (which are updated only once, at the very first key).

In general I think when there are multiple keys, the very first key should be the referencing data.