follow up of #1992

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

Author: wyli

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,
@@ -56,6 +57,10 @@ def __init__(
             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.
             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 = ensure_tuple_rep(meta_key_postfix, len(output_keys))
         self.postfix = postfix
         self.nearest_interp = ensure_tuple_rep(nearest_interp, len(self.output_keys))
         self._totensor = ToTensor()
@@ -90,7 +96,9 @@ def __call__(self, engine: Engine) -> None:
         Args:
             engine: Ignite Engine, it can be a trainer, validator or evaluator.
         """
-        for output_key, batch_key, nearest_interp in zip(self.output_keys, self.batch_keys, self.nearest_interp):
+        for output_key, batch_key, nearest_interp, meta_key in zip(
+            self.output_keys, self.batch_keys, self.nearest_interp, self.meta_key_postfix
+        ):
             transform_key = batch_key + InverseKeys.KEY_SUFFIX
             if transform_key not in engine.state.batch:
                 warnings.warn(f"all the transforms on `{batch_key}` are not InvertibleTransform.")
@@ -104,6 +112,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}_{meta_key}"
+            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}"

tests/test_handler_transform_inverter.py

@@ -32,6 +32,7 @@
     RandZoomd,
     ResizeWithPadOrCropd,
     ScaleIntensityd,
+    Spacingd,
     ToTensord,
 )
 from monai.utils.misc import set_determinism
@@ -48,6 +49,7 @@ def test_invert(self):
             [
                 LoadImaged(KEYS),
                 AddChanneld(KEYS),
+                Spacingd(KEYS, pixdim=(1.1, 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 +104,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 (0, 981), "diff. in two possible values")
 
 
 if __name__ == "__main__":