Repository: Project-MONAI/MONAI

Issue:
TestTimeAugmentation assumes label is present
**Describe the bug**
During training/testing/validation, ground truth labels are often present. However, at deployment only the image will be present. The way TTA is currently set up assumes that both image and label are present, which needs to be modified.

cc @diazandr3s 


Discussion:

Resolving patch:
diff --git a/monai/handlers/transform_inverter.py b/monai/handlers/transform_inverter.py
--- a/monai/handlers/transform_inverter.py
+++ b/monai/handlers/transform_inverter.py
@@ -87,9 +87,6 @@ def __init__(
                 each matches to the `output_keys` data.
             post_func: post processing for the inverted data, should be a callable function.
                 it also can be a list of callable, each matches to the `output_keys` data.
-            num_workers: number of workers when run data loader for inverse transforms,
-                default to 0 as only run one iteration and multi-processing may be even slower.
-                Set to `None`, to use the `num_workers` of the input transform data loader.
 
         """
         self.inverter = Invertd(
@@ -103,7 +100,6 @@ def __init__(
             to_tensor=to_tensor,
             device=device,
             post_func=post_func,
-            num_workers=num_workers,
         )
         self.output_keys = ensure_tuple(output_keys)
         self.meta_keys = ensure_tuple_rep(None, len(self.output_keys)) if meta_keys is None else ensure_tuple(meta_keys)
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
--- a/monai/transforms/post/dictionary.py
+++ b/monai/transforms/post/dictionary.py
@@ -429,7 +429,6 @@ def __init__(
         to_tensor: Union[bool, Sequence[bool]] = True,
         device: Union[Union[str, torch.device], Sequence[Union[str, torch.device]]] = "cpu",
         post_func: Union[Callable, Sequence[Callable]] = lambda x: x,
-        num_workers: Optional[int] = 0,
         allow_missing_keys: bool = False,
     ) -> None:
         """
@@ -465,9 +464,6 @@ def __init__(
                 each matches to the `keys` data.
             post_func: post processing for the inverted data, should be a callable function.
                 it also can be a list of callable, each matches to the `keys` data.
-            num_workers: number of workers when run data loader for inverse transforms,
-                default to 0 as only run one iteration and multi-processing may be even slower.
-                Set to `None`, to use the `num_workers` of the input transform data loader.
             allow_missing_keys: don't raise exception if key is missing.
 
         """


Test patch:
diff --git a/monai/data/test_time_augmentation.py b/monai/data/test_time_augmentation.py
--- a/monai/data/test_time_augmentation.py
+++ b/monai/data/test_time_augmentation.py
@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -21,7 +22,7 @@
 from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.inverse_batch_transform import BatchInverseTransform
 from monai.transforms.transform import Randomizable
-from monai.transforms.utils import allow_missing_keys_mode
+from monai.transforms.utils import allow_missing_keys_mode, convert_inverse_interp_mode
 from monai.utils.enums import CommonKeys, InverseKeys
 from monai.utils.module import optional_import
 
@@ -61,11 +62,11 @@ class TestTimeAugmentation:
         inferrer_fn: function to use to perform inference.
         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_keys: explicitly indicate the key of the expected meta data dictionary.
-            for example, for data with key `label`, the metadata by default is in `label_meta_dict`.
+        orig_key: the key of the original input data in the dict. will get the applied transform information
+            for this input data, then invert them for the expected data with `image_key`.
+        orig_meta_keys: the key of the meta data of original input data, will get the `affine`, `data_shape`, etc.
             the meta data is a dictionary object which contains: filename, original_shape, etc.
-            if None, will try to construct meta_keys by `key_{meta_key_postfix}`.
+            if None, will try to construct meta_keys by `{orig_key}_{meta_key_postfix}`.
         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
@@ -96,8 +97,9 @@ def __init__(
         inferrer_fn: Callable,
         device: Union[str, torch.device] = "cpu",
         image_key=CommonKeys.IMAGE,
-        label_key=CommonKeys.LABEL,
-        meta_keys: Optional[str] = None,
+        orig_key=CommonKeys.LABEL,
+        nearest_interp: bool = True,
+        orig_meta_keys: Optional[str] = None,
         meta_key_postfix="meta_dict",
         return_full_data: bool = False,
         progress: bool = True,
@@ -108,8 +110,9 @@ def __init__(
         self.inferrer_fn = inferrer_fn
         self.device = device
         self.image_key = image_key
-        self.label_key = label_key
-        self.meta_keys = meta_keys
+        self.orig_key = orig_key
+        self.nearest_interp = nearest_interp
+        self.orig_meta_keys = orig_meta_keys
         self.meta_key_postfix = meta_key_postfix
         self.return_full_data = return_full_data
         self.progress = progress
@@ -160,7 +163,7 @@ def __call__(
         ds = Dataset(data_in, self.transform)
         dl = DataLoader(ds, self.num_workers, batch_size=self.batch_size, collate_fn=pad_list_data_collate)
 
-        label_transform_key = self.label_key + InverseKeys.KEY_SUFFIX
+        transform_key = self.orig_key + InverseKeys.KEY_SUFFIX
 
         # create inverter
         inverter = BatchInverseTransform(self.transform, dl, collate_fn=list_data_collate)
@@ -178,19 +181,27 @@ def __call__(
             if isinstance(batch_output, np.ndarray):
                 batch_output = torch.Tensor(batch_output)
 
+            transform_info = batch_data[transform_key]
+            if self.nearest_interp:
+                transform_info = convert_inverse_interp_mode(
+                    trans_info=deepcopy(transform_info),
+                    mode="nearest",
+                    align_corners=None,
+                )
+
             # 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]}
+            inferred_dict = {self.orig_key: batch_output, transform_key: transform_info}
             # if meta dict is present, add that too (required for some inverse transforms)
-            label_meta_dict_key = self.meta_keys or 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]
+            meta_dict_key = self.orig_meta_keys or f"{self.orig_key}_{self.meta_key_postfix}"
+            if meta_dict_key in batch_data:
+                inferred_dict[meta_dict_key] = batch_data[meta_dict_key]
 
-            # do inverse transformation (allow missing keys as only inverting label)
+            # do inverse transformation (allow missing keys as only inverting the orig_key)
             with allow_missing_keys_mode(self.transform):  # type: ignore
                 inv_batch = inverter(inferred_dict)
 
             # append
-            outputs.append(inv_batch[self.label_key])
+            outputs.append(inv_batch[self.orig_key])
 
         # output
         output: np.ndarray = np.concatenate(outputs)
diff --git a/tests/test_handler_transform_inverter.py b/tests/test_handler_transform_inverter.py
--- a/tests/test_handler_transform_inverter.py
+++ b/tests/test_handler_transform_inverter.py
@@ -95,7 +95,6 @@ def _train_func(engine, batch):
             nearest_interp=True,
             to_tensor=[True, False],
             device="cpu",
-            num_workers=0 if sys.platform == "darwin" or torch.cuda.is_available() else 2,
         ).attach(engine)
 
         # test different nearest interpolation values
@@ -108,7 +107,6 @@ def _train_func(engine, batch):
             meta_key_postfix="meta_dict",
             nearest_interp=[True, False],
             post_func=[lambda x: x + 10, lambda x: x],
-            num_workers=0 if sys.platform == "darwin" or torch.cuda.is_available() else 2,
         ).attach(engine)
 
         engine.run(loader, max_epochs=1)
diff --git a/tests/test_invertd.py b/tests/test_invertd.py
--- a/tests/test_invertd.py
+++ b/tests/test_invertd.py
@@ -84,7 +84,6 @@ def test_invert(self):
             nearest_interp=True,
             to_tensor=[True, False, False],
             device="cpu",
-            num_workers=0 if sys.platform == "darwin" or torch.cuda.is_available() else 2,
         )
 
         # execute 1 epoch
diff --git a/tests/test_testtimeaugmentation.py b/tests/test_testtimeaugmentation.py
--- a/tests/test_testtimeaugmentation.py
+++ b/tests/test_testtimeaugmentation.py
@@ -148,13 +148,13 @@ def test_single_transform(self):
 
     def test_image_no_label(self):
         transforms = RandFlipd(["image"], prob=1.0)
-        tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x, label_key="image")
+        tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x, orig_key="image")
         tta(self.get_data(1, (20, 20), include_label=False))
 
     @unittest.skipUnless(has_nib, "Requires nibabel")
     def test_requires_meta_dict(self):
         transforms = Compose([RandFlipd("image"), Spacingd("image", pixdim=1.0)])
-        tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x, label_key="image")
+        tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x, orig_key="image")
         tta(self.get_data(1, (20, 20), include_label=False))