Repository: Project-MONAI/MONAI

Issue:
Spacing transforms scale pixel dimensions or resample to specified pixel dimensions
Hello! 

Is the Spacing transform meant to resample a given volume to a specified the pixel dimensions? I have NiFTI images with different voxel dimensions and am trying to resample all of them to (2mm, 2mm, 2mm) using the following code. \

```python 

fpaths = ["img0.nii", "img1.nii"]
transform = Compose([
        AsChannelFirst(channel_dim=-1),
        Spacing(pixdim=(2, 2, 2), mode="bilinear", diagonal=True),
        ScaleIntensity(minv=0.0, maxv=1.0),
        ToTensor(),
    ])
check_ds = ImageDataset(image_files=fpaths, transform=transform)

for i, img in enumerate(check_ds):

    print(check_ds[i][0].shape)

```

```img0.nii``` has a pixel size of (2mm, 2mm, 2mm) to begin with so it should remain unchanged by the transform, but its dimensions were shrunk by a factor of 2 from (128, 128, 90) to (64, 64, 46) ```img1.nii``` has a pixel size of (2mm, 2mm, 3.27mm), its first two dimensions should remain unchanged whereas the third dimension should expand to account for the reduced pixel size. However, its dimensions were reduced from (128, 128, 47) to (64, 64, 24). It appears that the Spacing transform only scales the pixel dimensions by a factor of 2 instead of resampling them to the specified dimensions. Is this the normal behavior? 

I am using Mona 0.5.2 and you can find the aforementioned NiFTI files here (https://drive.google.com/drive/folders/1k__dds_31elP-zkelCNSguWIDrHPTDe5?usp=sharing)

Thank you for looking into this. 


Discussion:
Hi @wyli ,

Could you please help share some comments in this ticket?

Thanks in advance.

the spacing transform requires additional meta data from the image, please [follow this example](https://github.com/Project-MONAI/tutorials/blob/master/3d_segmentation/torch/unet_training_dict.py#L63-L71) to construct a dictionary-based pipeline. The `ImageDataset` provides a simplified data source without the meta data.
```python
fpaths = ["img0.nii", "img1.nii"]
dict_transforms = Compose(
    [
        LoadImaged(keys=["img"], dtype=np.float32),
        # AddChanneld(keys=[f"img{i}" for i in range(4)]),
        AsChannelFirstd(keys=["img"], channel_dim=-1),
        Spacingd(keys=["img"], pixdim=(2, 2, 2), mode="bilinear"),
        ScaleIntensityd(keys=["img"]),
        ToTensord(keys=["img"]),
    ]
)

data_lst = [{"img": fpath} for fpath in fpaths]
check_ds = Dataset(data=data_lst, transform=dict_transforms)

for i, img in enumerate(check_ds):

    print(check_ds[i]["imp"].shape)
```

The above snippet generates tensors of the expected dimensions. But ```img1.nii``` was transformed to a volume of NaN's. Any thoughts on this? Thank you. 

Update: When I tried to save the processed images as NiFTI by the following transforms,

```python
Compose( [
    LoadImaged(keys=["img"]),
    AsChannelFirstd(keys=["img"], channel_dim=-1),
    Spacingd(
        keys=["img"], pixdim=(2, 2, 2), diagonal=True, mode="bilinear"
    ),
    ScaleIntensityd(keys=["img"], minv=0, maxv=1),
    SaveImaged(
        keys=["img"],
        output_dir="./visualization",
        output_ext=".nii",
        resample=True,
    ),
  ])
```

the attached error message was received

Traceback (most recent call last):
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/transforms/transform.py", line 48, in apply_transform
    return transform(data)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/transforms/io/dictionary.py", line 231, in __call__
    self._saver(img=d[key], meta_data=meta_data)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/transforms/io/array.py", line 282, in __call__
    self.saver.save(img, meta_data)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/data/nifti_saver.py", line 157, in save
    output_dtype=self.output_dtype,
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/data/nifti_writer.py", line 150, in write_nifti
    spatial_size=output_spatial_shape_[: len(data.shape)],
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/networks/layers/spatial_transforms.py", line 539, in forward
    grid = nn.functional.affine_grid(theta=theta[:, :sr], size=list(dst_size), align_corners=self.align_corners)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/torch/nn/functional.py", line 3931, in affine_grid
    raise ValueError("Expected non-zero, positive output size. Got {}".format(size))
ValueError: Expected non-zero, positive output size. Got [1, 1, -32768, -32768, 12032]

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/transforms/transform.py", line 48, in apply_transform
    return transform(data)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/transforms/compose.py", line 144, in __call__
    input_ = apply_transform(_transform, input_)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/transforms/transform.py", line 71, in apply_transform
    raise RuntimeError(f"applying transform {transform}") from e
RuntimeError: applying transform <monai.transforms.io.dictionary.SaveImaged object at 0x145d578f5be0>

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "preprocess.py", line 23, in <module>
    for i, img in enumerate(check_ds):
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/data/dataset.py", line 93, in __getitem__
    return self._transform(index)
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/data/dataset.py", line 79, in _transform
    return apply_transform(self.transform, data_i) if self.transform is not None else data_i
  File "/home/sfan/anaconda3/envs/torch_monai/lib/python3.6/site-packages/monai/transforms/transform.py", line 71, in apply_transform
    raise RuntimeError(f"applying transform {transform}") from e
RuntimeError: applying transform <monai.transforms.compose.Compose object at 0x145d578f5c88>

Resolving patch:
diff --git a/monai/data/image_dataset.py b/monai/data/image_dataset.py
--- a/monai/data/image_dataset.py
+++ b/monai/data/image_dataset.py
@@ -26,7 +26,8 @@ class ImageDataset(Dataset, Randomizable):
     for the image and segmentation arrays separately.
     The difference between this dataset and `ArrayDataset` is that this dataset can apply transform chain to images
     and segs and return both the images and metadata, and no need to specify transform to load images from files.
-
+    For more information, please see the image_dataset demo in the MONAI tutorial repo,
+    https://github.com/Project-MONAI/tutorials/blob/master/modules/image_dataset.ipynb
     """
 
     def __init__(
@@ -37,6 +38,7 @@ def __init__(
         transform: Optional[Callable] = None,
         seg_transform: Optional[Callable] = None,
         image_only: bool = True,
+        transform_with_metadata: bool = False,
         dtype: DtypeLike = np.float32,
         reader: Optional[Union[ImageReader, str]] = None,
         *args,
@@ -53,6 +55,7 @@ def __init__(
             transform: transform to apply to image arrays
             seg_transform: transform to apply to segmentation arrays
             image_only: if True return only the image volume, otherwise, return image volume and the metadata
+            transform_with_metadata: if True, the metadata will be passed to the transforms whenever possible.
             dtype: if not None convert the loaded image to this data type
             reader: register reader to load image file and meta data, if None, will use the default readers.
                 If a string of reader name provided, will construct a reader object with the `*args` and `**kwargs`
@@ -76,7 +79,10 @@ def __init__(
         self.labels = labels
         self.transform = transform
         self.seg_transform = seg_transform
+        if image_only and transform_with_metadata:
+            raise ValueError("transform_with_metadata=True requires image_only=False.")
         self.image_only = image_only
+        self.transform_with_metadata = transform_with_metadata
         self.loader = LoadImage(reader, image_only, dtype, *args, **kwargs)
         self.set_random_state(seed=get_seed())
         self._seed = 0  # transform synchronization seed
@@ -89,10 +95,9 @@ def randomize(self, data: Optional[Any] = None) -> None:
 
     def __getitem__(self, index: int):
         self.randomize()
-        meta_data = None
-        seg = None
-        label = None
+        meta_data, seg_meta_data, seg, label = None, None, None, None
 
+        # load data and optionally meta
         if self.image_only:
             img = self.loader(self.image_files[index])
             if self.seg_files is not None:
@@ -100,29 +105,40 @@ def __getitem__(self, index: int):
         else:
             img, meta_data = self.loader(self.image_files[index])
             if self.seg_files is not None:
-                seg, _ = self.loader(self.seg_files[index])
-
-        if self.labels is not None:
-            label = self.labels[index]
+                seg, seg_meta_data = self.loader(self.seg_files[index])
 
+        # apply the transforms
         if self.transform is not None:
             if isinstance(self.transform, Randomizable):
                 self.transform.set_random_state(seed=self._seed)
-            img = apply_transform(self.transform, img)
-
-        data = [img]
+            img = apply_transform(
+                self.transform, (img, meta_data) if self.transform_with_metadata else img, map_items=False
+            )
+            if self.transform_with_metadata:
+                img, meta_data = img
 
         if self.seg_transform is not None:
             if isinstance(self.seg_transform, Randomizable):
                 self.seg_transform.set_random_state(seed=self._seed)
-            seg = apply_transform(self.seg_transform, seg)
+            seg = apply_transform(
+                self.seg_transform, (seg, seg_meta_data) if self.transform_with_metadata else seg, map_items=False
+            )
+            if self.transform_with_metadata:
+                seg, seg_meta_data = seg
+
+        if self.labels is not None:
+            label = self.labels[index]
 
+        # construct outputs
+        data = [img]
         if seg is not None:
             data.append(seg)
         if label is not None:
             data.append(label)
         if not self.image_only and meta_data is not None:
             data.append(meta_data)
+        if not self.image_only and seg_meta_data is not None:
+            data.append(seg_meta_data)
         if len(data) == 1:
             return data[0]
         # use tuple instead of list as the default collate_fn callback of MONAI DataLoader flattens nested lists


Test patch:
diff --git a/tests/test_image_dataset.py b/tests/test_image_dataset.py
--- a/tests/test_image_dataset.py
+++ b/tests/test_image_dataset.py
@@ -17,7 +17,7 @@
 import numpy as np
 
 from monai.data import ImageDataset
-from monai.transforms.transform import RandomizableTransform
+from monai.transforms import Compose, EnsureChannelFirst, RandAdjustContrast, RandomizableTransform, Spacing
 
 FILENAMES = ["test1.nii.gz", "test2.nii", "test3.nii.gz"]
 
@@ -35,7 +35,39 @@ def __call__(self, data):
         return data + self._a
 
 
+class _TestCompose(Compose):
+    def __call__(self, input_):
+        data, meta = input_
+        data = self.transforms[0](data, meta)  # ensure channel first
+        data, _, meta["affine"] = self.transforms[1](data, meta["affine"])  # spacing
+        if len(self.transforms) == 3:
+            return self.transforms[2](data), meta  # image contrast
+        return data, meta
+
+
 class TestImageDataset(unittest.TestCase):
+    def test_use_case(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            img_ = nib.Nifti1Image(np.random.randint(0, 2, size=(20, 20, 20)), np.eye(4))
+            seg_ = nib.Nifti1Image(np.random.randint(0, 2, size=(20, 20, 20)), np.eye(4))
+            img_name, seg_name = os.path.join(tempdir, "img.nii.gz"), os.path.join(tempdir, "seg.nii.gz")
+            nib.save(img_, img_name)
+            nib.save(seg_, seg_name)
+            img_list, seg_list = [img_name], [seg_name]
+
+            img_xform = _TestCompose([EnsureChannelFirst(), Spacing(pixdim=(1.5, 1.5, 3.0)), RandAdjustContrast()])
+            seg_xform = _TestCompose([EnsureChannelFirst(), Spacing(pixdim=(1.5, 1.5, 3.0), mode="nearest")])
+            img_dataset = ImageDataset(
+                image_files=img_list,
+                seg_files=seg_list,
+                transform=img_xform,
+                seg_transform=seg_xform,
+                image_only=False,
+                transform_with_metadata=True,
+            )
+            self.assertTupleEqual(img_dataset[0][0].shape, (1, 14, 14, 7))
+            self.assertTupleEqual(img_dataset[0][1].shape, (1, 14, 14, 7))
+
     def test_dataset(self):
         with tempfile.TemporaryDirectory() as tempdir:
             full_names, ref_data = [], []
@@ -94,28 +126,30 @@ def test_dataset(self):
                 image_only=False,
             )
             for d_tuple, ref in zip(dataset, ref_data):
-                img, seg, meta = d_tuple
+                img, seg, meta, seg_meta = d_tuple
                 np.testing.assert_allclose(img, ref + 1, atol=1e-3)
                 np.testing.assert_allclose(seg, ref + 2, atol=1e-3)
                 np.testing.assert_allclose(meta["original_affine"], np.eye(4), atol=1e-3)
+                np.testing.assert_allclose(seg_meta["original_affine"], np.eye(4), atol=1e-3)
 
             # loading image/label, with meta
             dataset = ImageDataset(
                 full_names, transform=lambda x: x + 1, seg_files=full_names, labels=[1, 2, 3], image_only=False
             )
             for idx, (d_tuple, ref) in enumerate(zip(dataset, ref_data)):
-                img, seg, label, meta = d_tuple
+                img, seg, label, meta, seg_meta = d_tuple
                 np.testing.assert_allclose(img, ref + 1, atol=1e-3)
                 np.testing.assert_allclose(seg, ref, atol=1e-3)
                 np.testing.assert_allclose(idx + 1, label)
                 np.testing.assert_allclose(meta["original_affine"], np.eye(4), atol=1e-3)
+                np.testing.assert_allclose(seg_meta["original_affine"], np.eye(4), atol=1e-3)
 
             # loading image/label, with sync. transform
             dataset = ImageDataset(
                 full_names, transform=RandTest(), seg_files=full_names, seg_transform=RandTest(), image_only=False
             )
             for d_tuple, ref in zip(dataset, ref_data):
-                img, seg, meta = d_tuple
+                img, seg, meta, seg_meta = d_tuple
                 np.testing.assert_allclose(img, seg, atol=1e-3)
                 self.assertTrue(not np.allclose(img, ref))
                 np.testing.assert_allclose(meta["original_affine"], np.eye(4), atol=1e-3)