include meta when inverting (#1997)

* follow up of #1992

Signed-off-by: Wenqi Li <wenqil@nvidia.com>

* update spacing inverse

Signed-off-by: Wenqi Li <wenqil@nvidia.com>

Author: wyli

.github/workflows/pythonapp.yml

@@ -285,7 +285,7 @@ jobs:
     - name: Run quick tests (GPU)
       run: |
         nvidia-smi
-        export LAUNCH_DELAY=$(python -c "import numpy; print(numpy.random.randint(30) * 5)")
+        export LAUNCH_DELAY=$(python -c "import numpy; print(numpy.random.randint(30) * 10)")
         echo "Sleep $LAUNCH_DELAY"
         sleep $LAUNCH_DELAY
         export CUDA_VISIBLE_DEVICES=$(coverage run -m tests.utils)

docs/source/installation.md

@@ -59,11 +59,11 @@ for the latest features:
 
 ### Option 1 (as a part of your system-wide module):
 ```bash
-pip install git+https://github.com/Project-MONAI/MONAI#egg=MONAI
+pip install git+https://github.com/Project-MONAI/MONAI#egg=monai
 ```
 or, to build with MONAI Cpp/CUDA extensions:
 ```bash
-BUILD_MONAI=1 pip install git+https://github.com/Project-MONAI/MONAI#egg=MONAI
+BUILD_MONAI=1 pip install git+https://github.com/Project-MONAI/MONAI#egg=monai
 ```
 this command will download and install the current master branch of [MONAI from
 GitHub](https://github.com/Project-MONAI/MONAI).

monai/data/test_time_augmentation.py

@@ -62,6 +62,10 @@ class TestTimeAugmentation:
         device: device on which to perform inference.
         image_key: key used to extract image from input dictionary.
         label_key: key used to extract label from input dictionary.
+        meta_key_postfix: use `key_{postfix}` to to fetch the meta data according to the key data,
+            default is `meta_dict`, the meta data is a dictionary object.
+            For example, to handle key `image`,  read/write affine matrices from the
+            metadata `image_meta_dict` dictionary's `affine` field.
         return_full_data: normally, metrics are returned (mode, mean, std, vvc). Setting this flag to `True` will return the
             full data. Dimensions will be same size as when passing a single image through `inferrer_fn`, with a dimension appended
             equal in size to `num_examples` (N), i.e., `[N,C,H,W,[D]]`.
@@ -88,6 +92,7 @@ def __init__(
         device: Optional[Union[str, torch.device]] = "cuda" if torch.cuda.is_available() else "cpu",
         image_key=CommonKeys.IMAGE,
         label_key=CommonKeys.LABEL,
+        meta_key_postfix="meta_dict",
         return_full_data: bool = False,
         progress: bool = True,
     ) -> None:
@@ -98,6 +103,7 @@ def __init__(
         self.device = device
         self.image_key = image_key
         self.label_key = label_key
+        self.meta_key_postfix = meta_key_postfix
         self.return_full_data = return_full_data
         self.progress = progress
 
@@ -168,7 +174,7 @@ def __call__(
             # create a dictionary containing the inferred batch and their transforms
             inferred_dict = {self.label_key: batch_output, label_transform_key: batch_data[label_transform_key]}
             # if meta dict is present, add that too (required for some inverse transforms)
-            label_meta_dict_key = self.label_key + "_meta_dict"
+            label_meta_dict_key = f"{self.label_key}_{self.meta_key_postfix}"
             if label_meta_dict_key in batch_data:
                 inferred_dict[label_meta_dict_key] = batch_data[label_meta_dict_key]
 

monai/handlers/transform_inverter.py

@@ -40,6 +40,7 @@ def __init__(
         loader: TorchDataLoader,
         output_keys: Union[str, Sequence[str]] = CommonKeys.PRED,
         batch_keys: Union[str, Sequence[str]] = CommonKeys.IMAGE,
+        meta_key_postfix: str = "meta_dict",
         collate_fn: Optional[Callable] = no_collation,
         postfix: str = "_inverted",
         nearest_interp: Union[bool, Sequence[bool]] = True,
@@ -49,13 +50,17 @@ def __init__(
         Args:
             transform: a callable data transform on input data.
             loader: data loader used to run transforms and generate the batch of data.
-            collate_fn: how to collate data after inverse transformations.
-                default won't do any collation, so the output will be a list of size batch size.
             output_keys: the key of expected data in `ignite.engine.output`, invert transforms on it.
                 it also can be a list of keys, will invert transform for each of them. Default to "pred".
             batch_keys: the key of input data in `ignite.engine.batch`. will get the applied transforms
                 for this input data, then invert them for the expected data with `output_keys`.
                 It can also be a list of keys, each matches to the `output_keys` data. default to "image".
+            meta_key_postfix: use `{batch_key}_{postfix}` to to fetch the meta data according to the key data,
+                default is `meta_dict`, the meta data is a dictionary object.
+                For example, to handle key `image`,  read/write affine matrices from the
+                metadata `image_meta_dict` dictionary's `affine` field.
+            collate_fn: how to collate data after inverse transformations.
+                default won't do any collation, so the output will be a list of size batch size.
             postfix: will save the inverted result into `ignite.engine.output` with key `{output_key}{postfix}`.
             nearest_interp: whether to use `nearest` interpolation mode when inverting the spatial transforms,
                 default to `True`. If `False`, use the same interpolation mode as the original transform.
@@ -74,6 +79,7 @@ def __init__(
         )
         self.output_keys = ensure_tuple(output_keys)
         self.batch_keys = ensure_tuple_rep(batch_keys, len(self.output_keys))
+        self.meta_key_postfix = meta_key_postfix
         self.postfix = postfix
         self.nearest_interp = ensure_tuple_rep(nearest_interp, len(self.output_keys))
         self._totensor = ToTensor()
@@ -104,6 +110,9 @@ def __call__(self, engine: Engine) -> None:
                 batch_key: engine.state.output[output_key].detach().cpu(),
                 transform_key: transform_info,
             }
+            meta_dict_key = f"{batch_key}_{self.meta_key_postfix}"
+            if meta_dict_key in engine.state.batch:
+                segs_dict[meta_dict_key] = engine.state.batch[meta_dict_key]
 
             with allow_missing_keys_mode(self.transform):  # type: ignore
                 inverted_key = f"{output_key}{self.postfix}"

monai/transforms/spatial/array.py

@@ -129,6 +129,7 @@ def __call__(
         padding_mode: Optional[Union[GridSamplePadMode, str]] = None,
         align_corners: Optional[bool] = None,
         dtype: DtypeLike = None,
+        output_spatial_shape: Optional[np.ndarray] = None,
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
         """
         Args:
@@ -145,6 +146,9 @@ def __call__(
             dtype: data type for resampling computation. Defaults to ``self.dtype``.
                 If None, use the data type of input data. To be compatible with other modules,
                 the output data type is always ``np.float32``.
+            output_spatial_shape: specify the shape of the output data_array. This is typically useful for
+                the inverse of `Spacingd` where sometimes we could not compute the exact shape due to the quantization
+                error with the affines.
 
         Raises:
             ValueError: When ``data_array`` has no spatial dimensions.
@@ -195,7 +199,7 @@ def __call__(
             # AffineTransform requires a batch dim
             torch.as_tensor(np.ascontiguousarray(data_array).astype(_dtype)).unsqueeze(0),
             torch.as_tensor(np.ascontiguousarray(transform).astype(_dtype)),
-            spatial_size=output_shape,
+            spatial_size=output_shape if output_spatial_shape is None else output_spatial_shape,
         )
         output_data = np.asarray(output_data.squeeze(0).detach().cpu().numpy(), dtype=np.float32)  # type: ignore
         new_affine = to_affine_nd(affine, new_affine)

monai/transforms/spatial/dictionary.py

@@ -201,6 +201,7 @@ def __call__(
             meta_data = d[meta_data_key]
             # resample array of each corresponding key
             # using affine fetched from d[affine_key]
+            original_spatial_shape = d[key].shape[1:]
             d[key], old_affine, new_affine = self.spacing_transform(
                 data_array=np.asarray(d[key]),
                 affine=meta_data["affine"],
@@ -219,6 +220,7 @@ def __call__(
                     "padding_mode": padding_mode.value if isinstance(padding_mode, Enum) else padding_mode,
                     "align_corners": align_corners if align_corners is not None else "none",
                 },
+                orig_size=original_spatial_shape,
             )
             # set the 'affine' key
             meta_data["affine"] = new_affine
@@ -239,6 +241,7 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
             mode = transform[InverseKeys.EXTRA_INFO]["mode"]
             padding_mode = transform[InverseKeys.EXTRA_INFO]["padding_mode"]
             align_corners = transform[InverseKeys.EXTRA_INFO]["align_corners"]
+            orig_size = transform[InverseKeys.ORIG_SIZE]
             orig_pixdim = np.sqrt(np.sum(np.square(old_affine), 0))[:-1]
             inverse_transform = Spacing(orig_pixdim, diagonal=self.spacing_transform.diagonal)
             # Apply inverse
@@ -249,6 +252,7 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
                 padding_mode=padding_mode,
                 align_corners=False if align_corners == "none" else align_corners,
                 dtype=dtype,
+                output_spatial_shape=orig_size,
             )
             meta_data["affine"] = new_affine
             # Remove the applied transform
@@ -1482,7 +1486,7 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
                 align_corners=None if align_corners == "none" else align_corners,
             )
             # Size might be out by 1 voxel so pad
-            d[key] = SpatialPad(transform[InverseKeys.ORIG_SIZE])(d[key])
+            d[key] = SpatialPad(transform[InverseKeys.ORIG_SIZE], mode="edge")(d[key])
             # Remove the applied transform
             self.pop_transform(d, key)
 
@@ -1607,7 +1611,7 @@ def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndar
                     align_corners=None if align_corners == "none" else align_corners,
                 )
                 # Size might be out by 1 voxel so pad
-                d[key] = SpatialPad(transform[InverseKeys.ORIG_SIZE])(d[key])
+                d[key] = SpatialPad(transform[InverseKeys.ORIG_SIZE], mode="edge")(d[key])
             # Remove the applied transform
             self.pop_transform(d, key)
 

tests/test_handler_transform_inverter.py

@@ -24,6 +24,7 @@
     CastToTyped,
     Compose,
     LoadImaged,
+    Orientationd,
     RandAffined,
     RandAxisFlipd,
     RandFlipd,
@@ -32,6 +33,7 @@
     RandZoomd,
     ResizeWithPadOrCropd,
     ScaleIntensityd,
+    Spacingd,
     ToTensord,
 )
 from monai.utils.misc import set_determinism
@@ -48,6 +50,8 @@ def test_invert(self):
             [
                 LoadImaged(KEYS),
                 AddChanneld(KEYS),
+                Orientationd(KEYS, "RPS"),
+                Spacingd(KEYS, pixdim=(1.2, 1.01, 0.9), mode=["bilinear", "nearest"], dtype=np.float32),
                 ScaleIntensityd("image", minv=1, maxv=10),
                 RandFlipd(KEYS, prob=0.5, spatial_axis=[1, 2]),
                 RandAxisFlipd(KEYS, prob=0.5),
@@ -102,7 +106,8 @@ def _train_func(engine, batch):
         reverted_name = engine.state.output["label_meta_dict"]["filename_or_obj"][-1]
         original_name = data[-1]["label"]
         self.assertEqual(reverted_name, original_name)
-        self.assertTrue((reverted.size - n_good) in (0, 23641), "diff. in two possible values")
+        print("invert diff", reverted.size - n_good)
+        self.assertTrue((reverted.size - n_good) in (25300, 1812), "diff. in two possible values")
 
 
 if __name__ == "__main__":